Fix tls + creating tls_stream (#1058)

tap/tcp_reader.go

@@ -26,7 +26,7 @@ type tcpReader struct {
 	data          []byte
 	progress      *api.ReadProgress
 	captureTime   time.Time
-	parent        api.TcpStream
+	parent        *tcpStream
 	packetsSeen   uint
 	extension     *api.Extension
 	emitter       api.Emitter
@@ -35,7 +35,7 @@ type tcpReader struct {
 	sync.Mutex
 }
 
-func NewTcpReader(msgQueue chan api.TcpReaderDataMsg, progress *api.ReadProgress, ident string, tcpId *api.TcpID, captureTime time.Time, parent api.TcpStream, isClient bool, isOutgoing bool, extension *api.Extension, emitter api.Emitter, counterPair *api.CounterPair, reqResMatcher api.RequestResponseMatcher) api.TcpReader {
+func NewTcpReader(msgQueue chan api.TcpReaderDataMsg, progress *api.ReadProgress, ident string, tcpId *api.TcpID, captureTime time.Time, parent *tcpStream, isClient bool, isOutgoing bool, extension *api.Extension, emitter api.Emitter, counterPair *api.CounterPair, reqResMatcher api.RequestResponseMatcher) *tcpReader {
 	return &tcpReader{
 		msgQueue:      msgQueue,
 		progress:      progress,

tap/tcp_stream.go

@@ -26,7 +26,7 @@ type tcpStream struct {
 	sync.Mutex
 }
 
-func NewTcpStream(isTapTarget bool, streamsMap api.TcpStreamMap, capture api.Capture) api.TcpStream {
+func NewTcpStream(isTapTarget bool, streamsMap api.TcpStreamMap, capture api.Capture) *tcpStream {
 	return &tcpStream{
 		isTapTarget:     isTapTarget,
 		protoIdentifier: &api.ProtoIdentifier{},

tap/tcp_stream_factory.go

@@ -61,15 +61,14 @@ func (factory *tcpStreamFactory) New(net, transport gopacket.Flow, tcpLayer *lay
 	stream := NewTcpStream(isTapTarget, factory.streamsMap, getPacketOrigin(ac))
 	reassemblyStream := NewTcpReassemblyStream(fmt.Sprintf("%s:%s", net, transport), tcpLayer, fsmOptions, stream)
 	if stream.GetIsTapTarget() {
-		_stream := stream.(*tcpStream)
+		stream.setId(factory.streamsMap.NextId())
-		_stream.setId(factory.streamsMap.NextId())
 		for i, extension := range extensions {
 			reqResMatcher := extension.Dissector.NewResponseRequestMatcher()
 			counterPair := &api.CounterPair{
 				Request:  0,
 				Response: 0,
 			}
-			_stream.addClient(
+			stream.addClient(
 				NewTcpReader(
 					make(chan api.TcpReaderDataMsg),
 					&api.ReadProgress{},
@@ -90,7 +89,7 @@ func (factory *tcpStreamFactory) New(net, transport gopacket.Flow, tcpLayer *lay
 					reqResMatcher,
 				),
 			)
-			_stream.addServer(
+			stream.addServer(
 				NewTcpReader(
 					make(chan api.TcpReaderDataMsg),
 					&api.ReadProgress{},
@@ -112,11 +111,11 @@ func (factory *tcpStreamFactory) New(net, transport gopacket.Flow, tcpLayer *lay
 				),
 			)
 
-			factory.streamsMap.Store(stream.(*tcpStream).getId(), stream)
+			factory.streamsMap.Store(stream.getId(), stream)
 
 			factory.wg.Add(2)
-			go _stream.getClient(i).(*tcpReader).run(filteringOptions, &factory.wg)
+			go stream.getClient(i).(*tcpReader).run(filteringOptions, &factory.wg)
-			go _stream.getServer(i).(*tcpReader).run(filteringOptions, &factory.wg)
+			go stream.getServer(i).(*tcpReader).run(filteringOptions, &factory.wg)
 		}
 	}
 	return reassemblyStream

tap/tlstapper/tls_poller.go

@@ -21,34 +21,25 @@ import (
 )
 
 type tlsPoller struct {
-	tls             *TlsTapper
+	tls            *TlsTapper
-	readers         map[string]api.TcpReader
+	readers        map[string]*tlsReader
-	closedReaders   chan string
+	closedReaders  chan string
-	reqResMatcher   api.RequestResponseMatcher
+	reqResMatcher  api.RequestResponseMatcher
-	chunksReader    *perf.Reader
+	chunksReader   *perf.Reader
-	extension       *api.Extension
+	extension      *api.Extension
-	procfs          string
+	procfs         string
-	pidToNamespace  sync.Map
+	pidToNamespace sync.Map
-	isClosed        bool
-	protoIdentifier *api.ProtoIdentifier
-	isTapTarget     bool
-	origin          api.Capture
-	createdAt       time.Time
 }
 
 func newTlsPoller(tls *TlsTapper, extension *api.Extension, procfs string) *tlsPoller {
 	return &tlsPoller{
-		tls:             tls,
+		tls:           tls,
-		readers:         make(map[string]api.TcpReader),
+		readers:       make(map[string]*tlsReader),
-		closedReaders:   make(chan string, 100),
+		closedReaders: make(chan string, 100),
-		reqResMatcher:   extension.Dissector.NewResponseRequestMatcher(),
+		reqResMatcher: extension.Dissector.NewResponseRequestMatcher(),
-		extension:       extension,
+		extension:     extension,
-		chunksReader:    nil,
+		chunksReader:  nil,
-		procfs:          procfs,
+		procfs:        procfs,
-		protoIdentifier: &api.ProtoIdentifier{},
-		isTapTarget:     true,
-		origin:          api.Ebpf,
-		createdAt:       time.Now(),
 	}
 }
 
@@ -135,24 +126,13 @@ func (p *tlsPoller) handleTlsChunk(chunk *tlsChunk, extension *api.Extension,
 	key := buildTlsKey(chunk, ip, port)
 	reader, exists := p.readers[key]
 
-	newReader := NewTlsReader(
-		key,
-		func(r *tlsReader) {
-			p.closeReader(key, r)
-		},
-		chunk.isRequest(),
-		p,
-	)
-
 	if !exists {
-		reader = p.startNewTlsReader(chunk, ip, port, key, extension, newReader, options)
+		reader = p.startNewTlsReader(chunk, ip, port, key, emitter, extension, options)
 		p.readers[key] = reader
 	}
 
-	tlsReader := reader.(*tlsReader)
+	reader.captureTime = time.Now()
-
+	reader.chunks <- chunk
-	tlsReader.setCaptureTime(time.Now())
-	tlsReader.sendChunk(chunk)
 
 	if os.Getenv("MIZU_VERBOSE_TLS_TAPPER") == "true" {
 		p.logTls(chunk, ip, port)
@@ -161,25 +141,46 @@ func (p *tlsPoller) handleTlsChunk(chunk *tlsChunk, extension *api.Extension,
 	return nil
 }
 
-func (p *tlsPoller) startNewTlsReader(chunk *tlsChunk, ip net.IP, port uint16, key string, extension *api.Extension,
+func (p *tlsPoller) startNewTlsReader(chunk *tlsChunk, ip net.IP, port uint16, key string,
-	reader api.TcpReader, options *api.TrafficFilteringOptions) api.TcpReader {
+	emitter api.Emitter, extension *api.Extension, options *api.TrafficFilteringOptions) *tlsReader {
 
 	tcpid := p.buildTcpId(chunk, ip, port)
 
-	tlsReader := reader.(*tlsReader)
+	doneHandler := func(r *tlsReader) {
-	tlsReader.setTcpID(&tcpid)
+		p.closeReader(key, r)
+	}
 
-	tlsReader.setEmitter(&tlsEmitter{
+	tlsEmitter := &tlsEmitter{
-		delegate:  reader.GetEmitter(),
+		delegate:  emitter,
 		namespace: p.getNamespace(chunk.Pid),
-	})
+	}
+
+	reader := &tlsReader{
+		key:           key,
+		chunks:        make(chan *tlsChunk, 1),
+		doneHandler:   doneHandler,
+		progress:      &api.ReadProgress{},
+		tcpID:         &tcpid,
+		isClient:      chunk.isRequest(),
+		captureTime:   time.Now(),
+		extension:     extension,
+		emitter:       tlsEmitter,
+		counterPair:   &api.CounterPair{},
+		reqResMatcher: p.reqResMatcher,
+	}
+
+	stream := &tlsStream{
+		reader:          reader,
+		protoIdentifier: &api.ProtoIdentifier{},
+	}
+
+	reader.parent = stream
 
 	go dissect(extension, reader, options)
 	return reader
 }
 
-func dissect(extension *api.Extension, reader api.TcpReader,
+func dissect(extension *api.Extension, reader *tlsReader, options *api.TrafficFilteringOptions) {
-	options *api.TrafficFilteringOptions) {
 	b := bufio.NewReader(reader)
 
 	err := extension.Dissector.Dissect(b, reader, options)
@@ -279,27 +280,3 @@ func (p *tlsPoller) logTls(chunk *tlsChunk, ip net.IP, port uint16) {
 		srcIp, srcPort, dstIp, dstPort,
 		chunk.Recorded, chunk.Len, chunk.Start, str, hex.EncodeToString(chunk.Data[0:chunk.Recorded]))
 }
-
-func (p *tlsPoller) SetProtocol(protocol *api.Protocol) {
-	// TODO: Implement
-}
-
-func (p *tlsPoller) GetOrigin() api.Capture {
-	return p.origin
-}
-
-func (p *tlsPoller) GetProtoIdentifier() *api.ProtoIdentifier {
-	return p.protoIdentifier
-}
-
-func (p *tlsPoller) GetReqResMatcher() api.RequestResponseMatcher {
-	return p.reqResMatcher
-}
-
-func (p *tlsPoller) GetIsTapTarget() bool {
-	return p.isTapTarget
-}
-
-func (p *tlsPoller) GetIsClosed() bool {
-	return p.isClosed
-}

tap/tlstapper/tls_reader.go

@@ -14,41 +14,15 @@ type tlsReader struct {
 	doneHandler   func(r *tlsReader)
 	progress      *api.ReadProgress
 	tcpID         *api.TcpID
-	isClosed      bool
 	isClient      bool
 	captureTime   time.Time
-	parent        api.TcpStream
 	extension     *api.Extension
 	emitter       api.Emitter
 	counterPair   *api.CounterPair
+	parent        api.TcpStream
 	reqResMatcher api.RequestResponseMatcher
 }
 
-func NewTlsReader(key string, doneHandler func(r *tlsReader), isClient bool, stream api.TcpStream) api.TcpReader {
-	return &tlsReader{
-		key:         key,
-		chunks:      make(chan *tlsChunk, 1),
-		doneHandler: doneHandler,
-		parent:      stream,
-	}
-}
-
-func (r *tlsReader) sendChunk(chunk *tlsChunk) {
-	r.chunks <- chunk
-}
-
-func (r *tlsReader) setTcpID(tcpID *api.TcpID) {
-	r.tcpID = tcpID
-}
-
-func (r *tlsReader) setCaptureTime(captureTime time.Time) {
-	r.captureTime = captureTime
-}
-
-func (r *tlsReader) setEmitter(emitter api.Emitter) {
-	r.emitter = emitter
-}
-
 func (r *tlsReader) Read(p []byte) (int, error) {
 	var chunk *tlsChunk
 
@@ -111,7 +85,7 @@ func (r *tlsReader) GetEmitter() api.Emitter {
 }
 
 func (r *tlsReader) GetIsClosed() bool {
-	return r.isClosed
+	return false
 }
 
 func (r *tlsReader) GetExtension() *api.Extension {

tap/tlstapper/tls_stream.go

@@ -0,0 +1,32 @@
+package tlstapper
+
+import "github.com/up9inc/mizu/tap/api"
+
+type tlsStream struct {
+	reader          *tlsReader
+	protoIdentifier *api.ProtoIdentifier
+}
+
+func (t *tlsStream) GetOrigin() api.Capture {
+	return api.Ebpf
+}
+
+func (t *tlsStream) GetProtoIdentifier() *api.ProtoIdentifier {
+	return t.protoIdentifier
+}
+
+func (t *tlsStream) SetProtocol(protocol *api.Protocol) {
+	t.protoIdentifier.Protocol = protocol
+}
+
+func (t *tlsStream) GetReqResMatcher() api.RequestResponseMatcher {
+	return t.reader.reqResMatcher
+}
+
+func (t *tlsStream) GetIsTapTarget() bool {
+	return true
+}
+
+func (t *tlsStream) GetIsClosed() bool {
+	return false
+}