Golang message.Header類代碼示例

// Scans provided byte slice for the next record separator and tries to
// populate provided header and message objects using the scanned data.
// Returns new starting index into the passed buffer, or ok == false if no
// message was able to be extracted.
//
// TODO this code duplicates what is in the StreamParser and should be
// removed.
func findMessage(buf []byte, header *message.Header, msg *[]byte) (pos int, ok bool) {
	pos = bytes.IndexByte(buf, message.RECORD_SEPARATOR)
	if pos != -1 {
		if len(buf)-pos > 1 {
			headerLength := int(buf[pos+1])
			headerEnd := pos + headerLength + 3 // recsep+len+header+unitsep
			if len(buf) >= headerEnd {
				if header.MessageLength != nil || DecodeHeader(buf[pos+2:headerEnd], header) {
					messageEnd := headerEnd + int(header.GetMessageLength())
					if len(buf) >= messageEnd {
						*msg = (*msg)[:messageEnd-headerEnd]
						copy(*msg, buf[headerEnd:messageEnd])
						pos = messageEnd
						ok = true
					} else {
						*msg = (*msg)[:0]
					}
				} else {
					pos, ok = findMessage(buf[pos+1:], header, msg)
				}
			}
		}
	} else {
		pos = len(buf)
	}
	return
}

// Returns true if the provided message is unsigned or has a valid signature
// from one of the provided signers.
func authenticateMessage(signers map[string]Signer, header *message.Header,
	msg []byte) bool {

	digest := header.GetHmac()
	if digest != nil {
		var key string
		signer := fmt.Sprintf("%s_%d", header.GetHmacSigner(),
			header.GetHmacKeyVersion())
		if s, ok := signers[signer]; ok {
			key = s.HmacKey
		} else {
			return false
		}

		var hm hash.Hash
		switch header.GetHmacHashFunction() {
		case message.Header_MD5:
			hm = hmac.New(md5.New, []byte(key))
		case message.Header_SHA1:
			hm = hmac.New(sha1.New, []byte(key))
		}
		hm.Write(msg)
		expectedDigest := hm.Sum(nil)
		if subtle.ConstantTimeCompare(digest, expectedDigest) != 1 {
			return false
		}
	}
	return true
}

//.........這裏部分代碼省略.........
			}
			oRunner, err := NewFORunner("stoppingOutput", output, commonFO,
				"StoppingResumeOutput", chanSize)
			c.Assume(err, gs.IsNil)
			oRunner.maker = maker

			close(oRunner.inChan) // signal the shutdown
			mockHelper.EXPECT().PipelineConfig().Return(pConfig)
			var wg sync.WaitGroup
			wg.Add(1)
			oRunner.Start(mockHelper, &wg) // no panic => success
			wg.Wait()
			c.Expect(stopresumerunTimes, gs.Equals, 3)
			c.Expect(len(stopresumeHolder), gs.Equals, 2)
			c.Expect(stopresumeHolder[1], gs.Equals, "woot")
			c.Expect(oRunner.retainPack, gs.IsNil)
		})

		c.Specify("can exit without causing shutdown", func() {
			commonFO.Retries = RetryOptions{MaxRetries: 0}
			oRunner, err := NewFORunner("stoppingOutput", output, commonFO, "StoppingOutput",
				chanSize)
			c.Assume(err, gs.IsNil)
			oRunner.canExit = true
			oRunner.maker = maker

			// This pack is for the sending of the terminated message
			pack := NewPipelinePack(pConfig.injectRecycleChan)
			pConfig.injectRecycleChan <- pack

			// Feed in a pack to the input so we can verify its been recycled
			// after stopping (no leaks)
			pack = NewPipelinePack(pConfig.inputRecycleChan)
			oRunner.inChan <- pack

			// This is code to emulate the router removing the Output, and
			// closing up the outputs inChan channel
			go func() {
				<-pConfig.Router().RemoveOutputMatcher()
				// We don't close the matcher inChan because its not
				// instantiated in the tests
				close(oRunner.inChan)
			}()

			mockHelper.EXPECT().PipelineConfig().Return(pConfig)
			var wg sync.WaitGroup
			wg.Add(1)
			oRunner.Start(mockHelper, &wg)
			wg.Wait()
			c.Expect(stopoutputTimes, gs.Equals, 1)
			p := <-pConfig.router.inChan
			c.Expect(p.Message.GetType(), gs.Equals, "heka.terminated")
			c.Expect(p.Message.GetLogger(), gs.Equals, "hekad")
			plugin, _ := p.Message.GetFieldValue("plugin")
			c.Expect(plugin, gs.Equals, "stoppingOutput")
			// This should be 1 because the inChan should be flushed
			// and packs should not be leaked
			c.Expect(len(pConfig.inputRecycleChan), gs.Equals, 1)
		})

		c.Specify("encodes a message", func() {
			oRunner, err := NewFORunner("stoppingOutput", output, commonFO, "StoppingOutput",
				chanSize)
			c.Assume(err, gs.IsNil)
			oRunner.encoder = new(_payloadEncoder)
			oRunner.maker = maker
			_pack.Message = ts.GetTestMessage()
			payload := "Test Payload"

			c.Specify("without framing", func() {
				result, err := oRunner.Encode(_pack)
				c.Expect(err, gs.IsNil)
				c.Expect(string(result), gs.Equals, payload)
			})

			c.Specify("with framing", func() {
				oRunner.SetUseFraming(true)
				result, err := oRunner.Encode(_pack)
				c.Expect(err, gs.IsNil)

				i := bytes.IndexByte(result, message.UNIT_SEPARATOR)
				c.Expect(i > 3, gs.IsTrue, -1)
				c.Expect(string(result[i+1:]), gs.Equals, payload)
				header := new(message.Header)
				ok, err := message.DecodeHeader(result[2:i+1], header)
				c.Expect(ok, gs.IsTrue)
				c.Expect(err, gs.IsNil)
				c.Expect(header.GetMessageLength(), gs.Equals, uint32(len(payload)))
			})

			c.Specify("with framing, ignore message", func() {
				oRunner.SetUseFraming(true)
				oRunner.encoder = new(_ignoreEncoder)
				result, err := oRunner.Encode(_pack)
				c.Expect(err, gs.IsNil)
				c.Expect(result == nil, gs.IsTrue)
			})
		})
	})
}

func OutputRunnerSpec(c gs.Context) {
	t := new(ts.SimpleT)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockHelper := NewMockPluginHelper(ctrl)
	pc := new(PipelineConfig)
	pluginGlobals := new(PluginGlobals)

	c.Specify("Runner restarts a plugin on the first time only", func() {
		pluginGlobals.Retries = RetryOptions{
			MaxDelay:   "1us",
			Delay:      "1us",
			MaxJitter:  "1us",
			MaxRetries: 1,
		}
		pw := NewPluginWrapper("stoppingOutput")
		pw.ConfigCreator = func() interface{} { return nil }
		pw.PluginCreator = func() interface{} { return new(StoppingOutput) }
		output := new(StoppingOutput)
		pc.outputWrappers = make(map[string]*PluginWrapper)
		pc.outputWrappers["stoppingOutput"] = pw
		oRunner := NewFORunner("stoppingOutput", output, pluginGlobals)
		var wg sync.WaitGroup
		cfgCall := mockHelper.EXPECT().PipelineConfig()
		cfgCall.Return(pc)
		wg.Add(1)
		oRunner.Start(mockHelper, &wg) // no panic => success
		wg.Wait()
		c.Expect(stopoutputTimes, gs.Equals, 2)
	})

	c.Specify("Runner restarts plugin and resumes feeding it", func() {
		pluginGlobals.Retries = RetryOptions{
			MaxDelay:   "1us",
			Delay:      "1us",
			MaxJitter:  "1us",
			MaxRetries: 4,
		}
		pw := NewPluginWrapper("stoppingresumeOutput")
		pw.ConfigCreator = func() interface{} { return nil }
		pw.PluginCreator = func() interface{} { return new(StopResumeOutput) }
		output := new(StopResumeOutput)
		pc.outputWrappers = make(map[string]*PluginWrapper)
		pc.outputWrappers["stoppingresumeOutput"] = pw
		oRunner := NewFORunner("stoppingresumeOutput", output, pluginGlobals)
		var wg sync.WaitGroup
		cfgCall := mockHelper.EXPECT().PipelineConfig()
		cfgCall.Return(pc)
		wg.Add(1)
		oRunner.Start(mockHelper, &wg) // no panic => success
		wg.Wait()
		c.Expect(stopresumerunTimes, gs.Equals, 3)
		c.Expect(len(stopresumeHolder), gs.Equals, 2)
		c.Expect(stopresumeHolder[1], gs.Equals, "woot")
		c.Expect(oRunner.retainPack, gs.IsNil)
	})

	c.Specify("Runner encodes a message", func() {
		output := new(StoppingOutput)
		or := NewFORunner("test", output, pluginGlobals)
		or.encoder = new(_payloadEncoder)
		_pack.Message = ts.GetTestMessage()
		payload := "Test Payload"

		c.Specify("without framing", func() {
			result, err := or.Encode(_pack)
			c.Expect(err, gs.IsNil)
			c.Expect(string(result), gs.Equals, payload)
		})

		c.Specify("with framing", func() {
			or.SetUseFraming(true)
			result, err := or.Encode(_pack)
			c.Expect(err, gs.IsNil)

			i := bytes.IndexByte(result, message.UNIT_SEPARATOR)
			c.Expect(i > 3, gs.IsTrue, -1)
			c.Expect(string(result[i+1:]), gs.Equals, payload)
			header := new(message.Header)
			ok := DecodeHeader(result[2:i+1], header)
			c.Expect(ok, gs.IsTrue)
			c.Expect(header.GetMessageLength(), gs.Equals, uint32(len(payload)))
		})
	})
}