加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
input_tcp.go 2.14 KB
一键复制 编辑 原始数据 按行查看 历史
thinkerou 提交于 2017-08-11 19:28 . format file
package main
import (
"bufio"
"bytes"
"crypto/tls"
"fmt"
"io"
"log"
"net"
"os"
)
// TCPInput used for internal communication
type TCPInput struct {
data chan []byte
listener net.Listener
address string
config *TCPInputConfig
}
type TCPInputConfig struct {
secure bool
certificatePath string
keyPath string
}
// NewTCPInput constructor for TCPInput, accepts address with port
func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) {
i = new(TCPInput)
i.data = make(chan []byte, 1000)
i.address = address
i.config = config
i.listen(address)
return
}
func (i *TCPInput) Read(data []byte) (int, error) {
buf := <-i.data
copy(data, buf)
return len(buf), nil
}
func (i *TCPInput) listen(address string) {
if i.config.secure {
cer, err := tls.LoadX509KeyPair(i.config.certificatePath, i.config.keyPath)
if err != nil {
log.Fatal("Error while loading --input-file certificate:", err)
}
config := &tls.Config{Certificates: []tls.Certificate{cer}}
listener, err := tls.Listen("tcp", address, config)
if err != nil {
log.Fatal("Can't start --input-tcp with secure connection:", err)
}
i.listener = listener
} else {
listener, err := net.Listen("tcp", address)
if err != nil {
log.Fatal("Can't start:", err)
}
i.listener = listener
}
go func() {
for {
conn, err := i.listener.Accept()
if err != nil {
log.Println("Error while Accept()", err)
continue
}
go i.handleConnection(conn)
}
}()
}
func (i *TCPInput) handleConnection(conn net.Conn) {
defer conn.Close()
payloadSeparatorAsBytes := []byte(payloadSeparator)
reader := bufio.NewReader(conn)
var buffer bytes.Buffer
for {
line, err := reader.ReadBytes('\n')
if err != nil {
if err != io.EOF {
fmt.Fprintln(os.Stderr, "Unexpected error in input tcp connection:", err)
}
break
}
if bytes.Equal(payloadSeparatorAsBytes[1:], line) {
asBytes := buffer.Bytes()
buffer.Reset()
newBuf := make([]byte, len(asBytes)-1)
copy(newBuf, asBytes)
i.data <- newBuf
} else {
buffer.Write(line)
}
}
}
func (i *TCPInput) String() string {
return "TCP input: " + i.address
}
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化