diff --git a/proxy/tun/handler.go b/proxy/tun/handler.go index 56fe5653a..709fc1ec6 100644 --- a/proxy/tun/handler.go +++ b/proxy/tun/handler.go @@ -37,6 +37,25 @@ type Handler struct { downlinkCounter stats.Counter } +type tunUDPStatsWriter struct { + writer buf.Writer + counter stats.Counter +} + +func (w *tunUDPStatsWriter) WriteMultiBuffer(mb buf.MultiBuffer) error { + for len(mb) > 0 { + remaining, packet := buf.SplitFirst(mb) + packetSize := packet.Len() + if err := w.writer.WriteMultiBuffer(buf.MultiBuffer{packet}); err != nil { + buf.ReleaseMulti(remaining) + return err + } + w.counter.Add(int64(packetSize)) + mb = remaining + } + return nil +} + // ConnectionHandler interface with the only method that stack is going to push new connections to type ConnectionHandler interface { HandleConnection(conn net.Conn, destination net.Destination) @@ -171,7 +190,8 @@ func (t *Handler) HandleConnection(conn net.Conn, destination net.Destination) { return } source := net.DestinationFromAddr(remote) - if t.uplinkCounter != nil || t.downlinkCounter != nil { + isUDP := destination.Network == net.Network_UDP + if !isUDP && (t.uplinkCounter != nil || t.downlinkCounter != nil) { conn = &stat.CounterConnection{ Connection: conn, ReadCounter: t.uplinkCounter, @@ -203,9 +223,18 @@ func (t *Handler) HandleConnection(conn net.Conn, destination net.Destination) { }) errors.LogInfo(ctx, "processing from ", source, " to ", destination) + reader := &buf.TimeoutWrapperReader{Reader: buf.NewReader(conn)} + writer := buf.NewWriter(conn) + if isUDP { + reader.Counter = t.uplinkCounter + if t.downlinkCounter != nil { + writer = &tunUDPStatsWriter{writer: writer, counter: t.downlinkCounter} + } + } + link := &transport.Link{ - Reader: &buf.TimeoutWrapperReader{Reader: buf.NewReader(conn)}, - Writer: buf.NewWriter(conn), + Reader: reader, + Writer: writer, } if err := t.dispatcher.DispatchLink(ctx, destination, link); err != nil { errors.LogError(ctx, errors.New("connection closed").Base(err))