Move to ListenUDP + use channel

Author: simsor

proxy/agent.go

@@ -10,32 +10,41 @@ import (
 )
 
 type agent struct {
-remote net.Addr
-sock net.PacketConn
-ws *websocket.Conn
+sock *net.UDPConn
+ws *websocket.Conn
+addr *net.UDPAddr
+
+running bool
+udpData chan []byte
 }
 
 func (a *agent) ws2sock() {
 log := logrus.WithFields(logrus.Fields{
-"remote": a.remote.String(),
+"remote": a.addr.String(),
 "func": "ws2sock",
 })
-defer a.ws.Close()
+defer func() {
+a.ws.Close()
+a.running = false
+}()
 
 for {
 t, p, err := a.ws.ReadMessage()
 if err != nil {
 log.WithField("err", err).Error("Could not read WebSocket message")
 return
 }
-log.Info("Got WS message")
+
+if logExchanges {
+log.Info("Got WS message")
+}
 
 if t != websocket.BinaryMessage {
 log.Error("Wrong message type")
 return
 }
 
-n, err := a.sock.WriteTo(p, a.remote)
+n, err := a.sock.WriteTo(p, a.addr)
 if err != nil {
 log.WithError(err).Error("Could not write to socket")
 return
@@ -48,8 +57,41 @@ func (a *agent) ws2sock() {
 }).Error("Did not write full message")
 return
 }
-log.Info("Successfully sent WS message to socket")
 
-fmt.Println(hex.Dump(p))
+if logExchanges {
+log.Info("Successfully sent WS message to socket")
+}
+
+if hexdumpPackets {
+fmt.Println(hex.Dump(p))
+}
+}
+}
+
+func (a *agent) sock2ws() {
+log := logrus.WithFields(logrus.Fields{
+"remote": a.addr.String(),
+"func": "sock2ws",
+})
+
+defer func() {
+a.ws.Close()
+a.running = false
+}()
+
+for p := range a.udpData {
+err := a.ws.WriteMessage(websocket.BinaryMessage, p)
+if err != nil {
+log.WithField("err", err).Error("Could not WriteMessage()")
+return
+}
+
+if logExchanges {
+log.Info("Successfully sent UDP message to WebSocket")
+}
+
+if hexdumpPackets {
+fmt.Println(hex.Dump(p))
+}
 }
 }

proxy/server.go

@@ -1,8 +1,6 @@
 package proxy
 
 import (
-"encoding/hex"
-"fmt"
 "net"
 "net/url"
 
@@ -16,7 +14,7 @@ type SocketServer struct {
 ListenPort int
 Destination string
 
-listener net.PacketConn
+listener *net.UDPConn
 agents map[string]*agent
 }
 
@@ -32,7 +30,11 @@ func New(listen, dest string) SocketServer {
 
 func (s *SocketServer) listen() error {
 // Socket server
-listener, err := net.ListenPacket("udp4", fmt.Sprintf("%s:%d", s.ListenAddress, s.ListenPort))
+addr := net.UDPAddr{
+Port: s.ListenPort,
+IP: net.ParseIP(s.ListenAddress),
+}
+listener, err := net.ListenUDP("udp4", &addr)
 if err != nil {
 return err
 }
@@ -51,26 +53,23 @@ func (s *SocketServer) Start() error {
 p := make([]byte, 65535)
 
 for {
-n, addr, err := s.listener.ReadFrom(p)
+n, addr, err := s.listener.ReadFromUDP(p)
 if err != nil {
 logrus.WithFields(logrus.Fields{
 "err": err,
-}).Error("Could not ReadFrom() a connection")
+}).Error("Could not ReadFromUDP() a connection, removing")
+a := s.getAgent(addr)
+
+a.ws.Close()
+a.running = false
 continue
 }
-//logrus.WithField("remote", addr.String()).Info("Accepted UDP connection")
+
+data := make([]byte, n)
+copy(data, p[:n])
 
 a := s.getAgent(addr)
-err = a.ws.WriteMessage(websocket.BinaryMessage, p[:n])
-if err != nil {
-logrus.WithFields(logrus.Fields{
-"err": err,
-"remote": addr.String(),
-}).Error("Could not WriteMessage()")
-continue
-}
-logrus.WithField("remote", addr.String()).Info("Successfully sent UDP message to WebSocket")
-fmt.Println(hex.Dump(p[:n]))
+a.udpData <- data
 }
 
 }
@@ -80,9 +79,9 @@ func (s *SocketServer) Close() error {
 return s.listener.Close()
 }
 
-func (s *SocketServer) getAgent(ip net.Addr) *agent {
-a, ok := s.agents[ip.String()]
-if ok {
+func (s *SocketServer) getAgent(addr *net.UDPAddr) *agent {
+a, ok := s.agents[addr.String()]
+if ok && a.running {
 return a
 }
 
@@ -94,12 +93,18 @@ func (s *SocketServer) getAgent(ip net.Addr) *agent {
 }
 
 a = &agent{
-ws: ws,
-sock: s.listener,
-remote: ip,
+ws: ws,
+sock: s.listener,
+addr: addr,
+running: true,
 }
+a.udpData = make(chan []byte, 10000)
 go a.ws2sock()
-s.agents[ip.String()] = a
+go a.sock2ws()
+s.agents[addr.String()] = a
 
+if logNewConnections {
+logrus.WithField("remote", addr.String()).Info("New client")
+}
 return a
 }