Golang streambuf.New函數代碼示例

func sendAck(transp *mockTransport, seq uint32) {
	buf := streambuf.New(nil)
	buf.WriteByte('2')
	buf.WriteByte('A')
	buf.WriteNetUint32(seq)
	transp.sendBytes(buf.Bytes())
}

func TestSimpleEvent(t *testing.T) {
	transp := newMockTransport()
	client := newClientTestDriver(newLumberjackClient(transp, 5*time.Second))

	event := common.MapStr{"name": "me", "line": 10}
	client.Publish([]common.MapStr{event})

	// receive window message
	buf := streambuf.New(nil)
	win, err := recvMessage(buf, transp)
	assert.Nil(t, err)

	// receive data message
	msg, err := recvMessage(buf, transp)
	assert.Nil(t, err)

	// send ack
	sendAck(transp, 1)

	// stop test driver
	transp.Close()
	client.Stop()

	// validate
	assert.NotNil(t, win)
	assert.NotNil(t, msg)
	assert.Equal(t, 1, len(msg.events))
	msg = msg.events[0]
	assert.Equal(t, "me", msg.doc["name"])
	assert.Equal(t, 10.0, msg.doc["line"])
}

func (msg *udpMessage) addDatagram(
	header *mcUdpHeader,
	data []byte,
) *streambuf.Buffer {
	if msg.isComplete {
		return nil
	}

	if msg.numDatagrams == 1 {
		msg.isComplete = true
		return streambuf.NewFixed(data)
	}

	if msg.count < msg.numDatagrams {
		if msg.datagrams[header.seqNumber] != nil {
			return nil
		}
		msg.datagrams[header.seqNumber] = data
		msg.count++
	}

	if msg.count < msg.numDatagrams {
		return nil
	}

	buffer := streambuf.New(nil)
	for _, payload := range msg.datagrams {
		buffer.Append(payload)
	}
	msg.isComplete = true
	msg.datagrams = nil
	buffer.Fix()
	return buffer
}

func TestLumberjackTCP(t *testing.T) {
	useTLS := false
	var serverErr error
	var win, data *message

	// create server with randomized port
	listener, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatalf("Failed to create test server")
	}

	// create lumberjack output client
	config := outputs.MothershipConfig{
		TLS:     &useTLS,
		Timeout: 1,
		Hosts:   []string{listener.Addr().String()},
	}
	output := newTestLumberjackOutput(t, config)

	// start server
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		buf := streambuf.New(nil)
		client, err := listener.Accept()
		if err != nil {
			serverErr = err
			return
		}
		win, err = sockReadMessage(buf, client)
		if err != nil {
			serverErr = err
			return
		}
		data, err = sockReadMessage(buf, client)
		if err != nil {
			serverErr = err
			return
		}
		serverErr = sockSendACK(client, 1)
		return
	}()

	// send event to server
	event := common.MapStr{"name": "me", "line": 10}
	output.PublishEvent(time.Now(), event)

	wg.Wait()
	listener.Close()

	// validate output
	assert.Nil(t, serverErr)
	assert.NotNil(t, win)
	assert.NotNil(t, data)
	assert.Equal(t, 1, len(data.events))
	data = data.events[0]
	assert.Equal(t, "\"me\"", data.kv["name"])
	assert.Equal(t, "10", data.kv["line"])
}

func sockSendACK(out io.Writer, seq uint32) error {
	buf := streambuf.New(nil)
	buf.WriteByte('2')
	buf.WriteByte('A')
	buf.WriteNetUint32(seq)
	_, err := out.Write(buf.Bytes())
	return err
}

func (l *lineReader) init(
	input io.Reader,
	codec encoding.Encoding,
	bufferSize int,
) error {
	l.rawInput = input
	l.codec = codec
	l.bufferSize = bufferSize

	l.codec.NewEncoder()
	nl, _, err := transform.Bytes(l.codec.NewEncoder(), []byte{'\n'})
	if err != nil {
		return err
	}

	l.nl = nl
	l.decoder = l.codec.NewDecoder()
	l.inBuffer = streambuf.New(nil)
	l.outBuffer = streambuf.New(nil)
	return nil
}

func TestStructuredEvent(t *testing.T) {
	transp := newMockTransport()
	client := newClientTestDriver(newLumberjackClient(transp, 5*time.Second))
	event := common.MapStr{
		"name": "test",
		"struct": common.MapStr{
			"field1": 1,
			"field2": true,
			"field3": []int{1, 2, 3},
			"field4": []interface{}{
				1,
				"test",
				common.MapStr{
					"sub": "field",
				},
			},
			"field5": common.MapStr{
				"sub1": 2,
			},
		},
	}
	client.Publish([]common.MapStr{event})

	buf := streambuf.New(nil)
	win, err := recvMessage(buf, transp)
	assert.Nil(t, err)

	msg, err := recvMessage(buf, transp)
	assert.Nil(t, err)

	sendAck(transp, 1)

	transp.Close()
	client.Stop()

	// validate
	assert.NotNil(t, win)
	assert.NotNil(t, msg)
	assert.Equal(t, 1, len(msg.events))
	msg = msg.events[0]
	assert.Equal(t, "\"test\"", msg.kv["name"])
	assert.Equal(t, "1", msg.kv["struct.field1"])
	assert.Equal(t, "true", msg.kv["struct.field2"])
	assert.Equal(t, "[1,2,3]", msg.kv["struct.field3"])
	assert.Equal(t, "[1,\"test\",{\"sub\":\"field\"}]", msg.kv["struct.field4"])
	assert.Equal(t, "2", msg.kv["struct.field5.sub1"])
}

func testConnectionType(
	t *testing.T,
	server *mockLSServer,
	makeOutputer func() outputs.BulkOutputer,
) {
	var result struct {
		err       error
		win, data *message
		signal    bool
	}

	var wg struct {
		ready  sync.WaitGroup
		finish sync.WaitGroup
	}

	wg.ready.Add(1)  // server signaling readiness to client worker
	wg.finish.Add(2) // server/client signaling test end

	// server loop
	go func() {
		defer wg.finish.Done()
		wg.ready.Done()

		client := server.accept()
		server.handshake(client)

		buf := streambuf.New(nil)
		result.win = server.readMessage(buf, client)
		result.data = server.readMessage(buf, client)
		server.sendACK(client, 1)
		result.err = server.err
	}()

	// worker loop
	go func() {
		defer wg.finish.Done()
		wg.ready.Wait()

		output := makeOutputer()

		signal := outputs.NewSyncSignal()
		output.PublishEvent(signal, time.Now(), testEvent())
		result.signal = signal.Wait()
	}()

	// wait shutdown
	wg.finish.Wait()
	server.Close()

	// validate output
	assert.Nil(t, result.err)
	assert.True(t, result.signal)

	data := result.data
	assert.NotNil(t, result.win)
	assert.NotNil(t, result.data)
	if data != nil {
		assert.Equal(t, 1, len(data.events))
		data = data.events[0]
		assert.Equal(t, 10.0, data.doc["extra"])
		assert.Equal(t, "message", data.doc["message"])
	}

}

func readMessage(buf *streambuf.Buffer) (*message, error) {
	if !buf.Avail(2) {
		return nil, nil
	}

	version, _ := buf.ReadNetUint8At(0)
	if version != '2' {
		return nil, errors.New("version error")
	}

	code, _ := buf.ReadNetUint8At(1)
	switch code {
	case 'W':
		if !buf.Avail(6) {
			return nil, nil
		}
		size, _ := buf.ReadNetUint32At(2)
		buf.Advance(6)
		buf.Reset()
		return &message{code: code, size: size}, buf.Err()
	case 'C':
		if !buf.Avail(6) {
			return nil, nil
		}
		len, _ := buf.ReadNetUint32At(2)
		if !buf.Avail(int(len) + 6) {
			return nil, nil
		}
		buf.Advance(6)

		tmp, _ := buf.Collect(int(len))
		buf.Reset()

		dataBuf := streambuf.New(nil)
		// decompress data
		decomp, err := zlib.NewReader(streambuf.NewFixed(tmp))
		if err != nil {
			return nil, err
		}
		// dataBuf.ReadFrom(streambuf.NewFixed(tmp))
		dataBuf.ReadFrom(decomp)
		decomp.Close()

		// unpack data
		dataBuf.Fix()
		var events []*message
		for dataBuf.Len() > 0 {
			version, _ := dataBuf.ReadNetUint8()
			if version != '2' {
				return nil, errors.New("version error 2")
			}

			code, _ := dataBuf.ReadNetUint8()
			if code != 'J' {
				return nil, errors.New("expected json data frame")
			}

			seq, _ := dataBuf.ReadNetUint32()
			payloadLen, _ := dataBuf.ReadNetUint32()
			jsonRaw, _ := dataBuf.Collect(int(payloadLen))

			var doc interface{}
			err = json.Unmarshal(jsonRaw, &doc)
			if err != nil {
				return nil, err
			}

			events = append(events, &message{
				code: code,
				seq:  seq,
				doc:  doc.(map[string]interface{}),
			})
		}
		return &message{code: 'C', events: events}, nil
	default:
		return nil, errors.New("unknown code")
	}
}

func TestLogstashTLS(t *testing.T) {
	pem := "ca_test.pem"
	key := "ca_test.key"

	var serverErr error
	var win, data *message

	genCertsIfMIssing(t, pem, key)

	tcpListener, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatalf("failed to generate TCP listener")
	}

	// create lumberjack output client
	config := outputs.MothershipConfig{
		TLS: &outputs.TLSConfig{
			Certificate:    pem,
			CertificateKey: key,
			CAs:            []string{pem},
		},
		Timeout: 5,
		Hosts:   []string{tcpListener.Addr().String()},
	}

	tlsConfig, err := outputs.LoadTLSConfig(config.TLS)
	if err != nil {
		tcpListener.Close()
		t.Fatalf("failed to load certificates")
	}

	listener := tls.NewListener(tcpListener, tlsConfig)

	// start server
	var wg sync.WaitGroup
	var wgReady sync.WaitGroup
	wg.Add(1)
	wgReady.Add(1)
	go func() {
		defer wg.Done()

		for i := 0; i < 3; i++ { // try up to 3 failed connection attempts
			// server read timeout
			timeout := 5 * time.Second
			buf := streambuf.New(nil)
			wgReady.Done()
			client, err := listener.Accept()
			if err != nil {
				continue
			}

			tlsConn, ok := client.(*tls.Conn)
			if !ok {
				serverErr = errors.New("no tls connection")
				return
			}

			if err := client.SetDeadline(time.Now().Add(timeout)); err != nil {
				serverErr = err
				return
			}

			err = tlsConn.Handshake()
			if err != nil {
				serverErr = err
				return
			}

			if err := client.SetDeadline(time.Now().Add(timeout)); err != nil {
				serverErr = err
				return
			}
			win, err = sockReadMessage(buf, client)
			if err != nil {
				serverErr = err
				return
			}

			if err := client.SetDeadline(time.Now().Add(timeout)); err != nil {
				serverErr = err
				return
			}
			data, err = sockReadMessage(buf, client)
			if err != nil {
				serverErr = err
				return
			}
			serverErr = sockSendACK(client, 1)

			return
		}
	}()

	// send event to server
	go func() {
		wgReady.Wait()
		output := newTestLumberjackOutput(t, "", &config)

		event := testEvent()
		output.PublishEvent(nil, time.Now(), event)
//.........這裏部分代碼省略.........

func TestLogstashTCP(t *testing.T) {
	var serverErr error
	var win, data *message

	// create server with randomized port
	listener, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatalf("Failed to create test server")
	}

	// start server
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()

		// server read timeout
		timeout := 5 * time.Second
		buf := streambuf.New(nil)

		client, err := listener.Accept()
		if err != nil {
			t.Logf("failed on accept: %v", err)
			serverErr = err
			return
		}

		if err := client.SetDeadline(time.Now().Add(timeout)); err != nil {
			serverErr = err
			return
		}
		win, err = sockReadMessage(buf, client)
		if err != nil {
			t.Logf("failed on read window size: %v", err)
			serverErr = err
			return
		}

		if err := client.SetDeadline(time.Now().Add(timeout)); err != nil {
			serverErr = err
			return
		}
		data, err = sockReadMessage(buf, client)
		if err != nil {
			t.Logf("failed on read data frame: %v", err)
			serverErr = err
			return
		}

		if err := client.SetDeadline(time.Now().Add(timeout)); err != nil {
			serverErr = err
			return
		}
		err = sockSendACK(client, 1)
		if err != nil {
			t.Logf("failed on read data frame: %v", err)
			serverErr = err
		}
	}()

	// create lumberjack output client
	config := outputs.MothershipConfig{
		TLS:     nil,
		Timeout: 2,
		Hosts:   []string{listener.Addr().String()},
	}
	output := newTestLumberjackOutput(t, "", &config)

	// send event to server
	event := testEvent()
	output.PublishEvent(nil, time.Now(), event)

	wg.Wait()
	listener.Close()

	// validate output
	assert.Nil(t, serverErr)
	assert.NotNil(t, win)
	if data == nil {
		t.Fatalf("No data received")
	}
	assert.Equal(t, 1, len(data.events))
	data = data.events[0]
	assert.Equal(t, 10.0, data.doc["extra"])
	assert.Equal(t, "message", data.doc["message"])
}

func readMessage(buf *streambuf.Buffer) (*message, error) {
	if !buf.Avail(2) {
		return nil, nil
	}

	version, _ := buf.ReadNetUint8At(0)
	if version != '1' {
		return nil, errors.New("version error")
	}

	code, _ := buf.ReadNetUint8At(1)
	switch code {
	case 'W':
		if !buf.Avail(6) {
			return nil, nil
		}
		size, _ := buf.ReadNetUint32At(2)
		buf.Advance(6)
		buf.Reset()
		return &message{code: code, size: size}, buf.Err()
	case 'C':
		if !buf.Avail(6) {
			return nil, nil
		}
		len, _ := buf.ReadNetUint32At(2)
		if !buf.Avail(int(len) + 6) {
			return nil, nil
		}
		buf.Advance(6)

		tmp, _ := buf.Collect(int(len))
		buf.Reset()

		dataBuf := streambuf.New(nil)
		// decompress data
		decomp, err := zlib.NewReader(streambuf.NewFixed(tmp))
		if err != nil {
			return nil, err
		}
		// dataBuf.ReadFrom(streambuf.NewFixed(tmp))
		dataBuf.ReadFrom(decomp)
		decomp.Close()

		// unpack data
		dataBuf.Fix()
		var events []*message
		for dataBuf.Len() > 0 {
			version, _ := dataBuf.ReadNetUint8()
			if version != '1' {
				return nil, errors.New("version error 2")
			}

			code, _ := dataBuf.ReadNetUint8()
			if code != 'D' {
				return nil, errors.New("expected data frame")
			}

			seq, _ := dataBuf.ReadNetUint32()
			pairCount, _ := dataBuf.ReadNetUint32()
			kv := make(map[string]string)
			for i := 0; i < int(pairCount); i++ {
				keyLen, _ := dataBuf.ReadNetUint32()
				keyRaw, _ := dataBuf.Collect(int(keyLen))
				valLen, _ := dataBuf.ReadNetUint32()
				valRaw, _ := dataBuf.Collect(int(valLen))
				kv[string(keyRaw)] = string(valRaw)
			}

			events = append(events, &message{code: code, seq: seq, kv: kv})
		}
		return &message{code: 'C', events: events}, nil
	default:
		return nil, errors.New("unknown code")
	}
}

func TestLumberjackTLS(t *testing.T) {
	pem := "ca_test.pem"
	key := "ca_test.key"

	var serverErr error
	var win, data *message

	genCertsIfMIssing(t, pem, key)

	tlsFiles, err := loadTLSConfig(&TLSConfig{
		Certificate: pem,
		Key:         key,
		CAs:         []string{pem},
	})
	if err != nil {
		t.Fatalf("failed to load certificates")
	}

	var tlsconfig tls.Config
	tlsconfig.MinVersion = tls.VersionTLS10
	tlsconfig.RootCAs = tlsFiles.RootCAs
	tlsconfig.Certificates = tlsFiles.Certificates
	tlsconfig.ServerName = "localhost"

	tcpListener, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatalf("failed to generate TCP listener")
	}

	listener := tls.NewListener(tcpListener, &tlsconfig)

	// create lumberjack output client
	useTLS := true
	config := outputs.MothershipConfig{
		TLS:            &useTLS,
		Timeout:        5,
		Hosts:          []string{listener.Addr().String()},
		Certificate:    pem,
		CertificateKey: key,
		CAs:            []string{pem},
	}

	// start server
	var wg sync.WaitGroup
	var wgReady sync.WaitGroup
	wg.Add(1)
	wgReady.Add(1)
	go func() {
		defer wg.Done()

		for i := 0; i < 3; i++ { // try up to 3 failed connection attempts
			buf := streambuf.New(nil)
			wgReady.Done()
			client, err := listener.Accept()
			if err != nil {
				continue
			}

			tlsConn, ok := client.(*tls.Conn)
			if !ok {
				serverErr = errors.New("no tls connection")
				return
			}

			err = tlsConn.Handshake()
			if err != nil {
				serverErr = err
				return
			}

			win, err = sockReadMessage(buf, client)
			if err != nil {
				serverErr = err
				return
			}

			data, err = sockReadMessage(buf, client)
			if err != nil {
				serverErr = err
				return
			}
			serverErr = sockSendACK(client, 1)

			return
		}
	}()

	// send event to server
	go func() {
		wgReady.Wait()
		output := newTestLumberjackOutput(t, config)

		event := common.MapStr{"name": "me", "line": 10}
		output.PublishEvent(time.Now(), event)
	}()

	wg.Wait()
	listener.Close()

	// validate output
//.........這裏部分代碼省略.........