let parent process clean up

Author: glaslos

app/server.go

@@ -7,10 +7,10 @@ import (
 	"os"
 	"os/signal"
 	"runtime/debug"
-	"sync"
 	"syscall"
 
 	"github.com/mushorg/glutton"
+
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 )
@@ -32,7 +32,7 @@ func main() {
  \_____|_|\__,_|\__|\__\___/|_| |_|
 
 	`)
-	fmt.Printf("%s %s\n", VERSION, BUILDDATE)
+	fmt.Printf("%s %s\n\n", VERSION, BUILDDATE)
 
 	pflag.StringP("interface", "i", "eth0", "Bind to this interface")
 	pflag.IntP("ssh", "s", 0, "Override SSH port")
@@ -53,39 +53,37 @@ func main() {
 		return
 	}
 
-	gtn, err := glutton.New(context.Background())
+	g, err := glutton.New(context.Background())
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	if err := gtn.Init(); err != nil {
+	if err := g.Init(); err != nil {
 		log.Fatal("Failed to initialize Glutton:", err)
 	}
 
-	exitMtx := sync.RWMutex{}
 	exit := func() {
 		// See if there was a panic...
 		if r := recover(); r != nil {
 			fmt.Fprintln(os.Stderr, r)
 			fmt.Println("stacktrace from panic: \n" + string(debug.Stack()))
 		}
-		exitMtx.Lock()
-		gtn.Shutdown()
-		exitMtx.Unlock()
+		g.Shutdown()
 	}
-	defer exit()
 
 	// capture and handle signals
 	sig := make(chan os.Signal, 1)
 	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
 	go func() {
 		<-sig
+		fmt.Print("\r")
 		exit()
-		fmt.Println("\nleaving...")
+		fmt.Println()
 		os.Exit(0)
 	}()
 
-	if err := gtn.Start(); err != nil {
+	if err := g.Start(); err != nil {
+		exit()
 		log.Fatal("Failed to start Glutton server:", err)
 	}
 }

glutton.go

@@ -148,55 +148,71 @@ func (g *Glutton) Init() error {
 }
 
 func (g *Glutton) udpListen(wg *sync.WaitGroup) {
-	defer wg.Done()
+	defer func() {
+		wg.Done()
+	}()
 	buffer := make([]byte, 1024)
 	for {
-		n, srcAddr, dstAddr, err := tproxy.ReadFromUDP(g.Server.udpListener, buffer)
+		select {
+		case <-g.ctx.Done():
+			if err := g.Server.udpConn.Close(); err != nil {
+				g.Logger.Error("Failed to close UDP listener", producer.ErrAttr(err))
+			}
+			return
+		default:
+		}
+		g.Server.udpConn.SetReadDeadline(time.Now().Add(1 * time.Second))
+		n, srcAddr, dstAddr, err := tproxy.ReadFromUDP(g.Server.udpConn, buffer)
 		if err != nil {
-			g.Logger.Error("failed to read UDP packet", producer.ErrAttr(err))
+			if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
+				continue
+			}
+			g.Logger.Error("Failed to read UDP packet", producer.ErrAttr(err))
 		}
 
 		rule, err := g.applyRules("udp", srcAddr, dstAddr)
 		if err != nil {
-			g.Logger.Error("failed to apply rules", producer.ErrAttr(err))
+			g.Logger.Error("Failed to apply rules", producer.ErrAttr(err))
 		}
 		if rule == nil {
 			rule = &rules.Rule{Target: "udp"}
 		}
 		md, err := g.connTable.Register(srcAddr.IP.String(), strconv.Itoa(int(srcAddr.AddrPort().Port())), dstAddr.AddrPort().Port(), rule)
 		if err != nil {
-			g.Logger.Error("failed to register UDP packet", producer.ErrAttr(err))
+			g.Logger.Error("Failed to register UDP packet", producer.ErrAttr(err))
 		}
 
 		if hfunc, ok := g.udpProtocolHandlers[rule.Target]; ok {
 			data := buffer[:n]
 			go func() {
 				if err := hfunc(g.ctx, srcAddr, dstAddr, data, md); err != nil {
-					g.Logger.Error("failed to handle UDP payload", producer.ErrAttr(err))
+					g.Logger.Error("Failed to handle UDP payload", producer.ErrAttr(err))
 				}
 			}()
 		}
 	}
 }
 
-func (g *Glutton) tcpListen(wg *sync.WaitGroup) {
-	defer wg.Done()
+func (g *Glutton) tcpListen() {
 	for {
 		select {
 		case <-g.ctx.Done():
+			if err := g.Server.tcpListener.Close(); err != nil {
+				g.Logger.Error("Failed to close TCP listener", producer.ErrAttr(err))
+			}
 			return
 		default:
 		}
 
 		conn, err := g.Server.tcpListener.Accept()
 		if err != nil {
-			g.Logger.Error("failed to accept connection", producer.ErrAttr(err))
+			g.Logger.Error("Failed to accept connection", producer.ErrAttr(err))
 			continue
 		}
 
 		rule, err := g.applyRulesOnConn(conn)
 		if err != nil {
-			g.Logger.Error("failed to apply rules", producer.ErrAttr(err))
+			g.Logger.Error("Failed to apply rules", producer.ErrAttr(err))
 			continue
 		}
 		if rule == nil {
@@ -205,21 +221,21 @@ func (g *Glutton) tcpListen(wg *sync.WaitGroup) {
 
 		md, err := g.connTable.RegisterConn(conn, rule)
 		if err != nil {
-			g.Logger.Error("failed to register connection", producer.ErrAttr(err))
+			g.Logger.Error("Failed to register connection", producer.ErrAttr(err))
 			continue
 		}
 
 		g.Logger.Debug("new connection", slog.String("addr", conn.LocalAddr().String()), slog.String("handler", rule.Target))
 
 		g.ctx = context.WithValue(g.ctx, ctxTimeout("timeout"), int64(viper.GetInt("conn_timeout")))
 		if err := g.UpdateConnectionTimeout(g.ctx, conn); err != nil {
-			g.Logger.Error("failed to set connection timeout", producer.ErrAttr(err))
+			g.Logger.Error("Failed to set connection timeout", producer.ErrAttr(err))
 		}
 
 		if hfunc, ok := g.tcpProtocolHandlers[rule.Target]; ok {
 			go func() {
 				if err := hfunc(g.ctx, conn, md); err != nil {
-					g.Logger.Error("failed to handle TCP connection", producer.ErrAttr(err), slog.String("handler", rule.Target))
+					g.Logger.Error("Failed to handle TCP connection", producer.ErrAttr(err), slog.String("handler", rule.Target))
 				}
 			}()
 		}
@@ -228,13 +244,7 @@ func (g *Glutton) tcpListen(wg *sync.WaitGroup) {
 
 // Start the listener, this blocks for new connections
 func (g *Glutton) Start() error {
-	quit := make(chan struct{}) // stop monitor on shutdown
-	defer func() {
-		quit <- struct{}{}
-		g.Shutdown()
-	}()
-
-	g.startMonitor(quit)
+	g.startMonitor()
 
 	sshPort := viper.GetUint32("ports.ssh")
 	if err := setTProxyIPTables(viper.GetString("interface"), g.publicAddrs[0].String(), "tcp", uint32(g.Server.tcpPort), sshPort); err != nil {
@@ -249,9 +259,7 @@ func (g *Glutton) Start() error {
 
 	wg.Add(1)
 	go g.udpListen(wg)
-
-	wg.Add(1)
-	go g.tcpListen(wg)
+	go g.tcpListen()
 
 	wg.Wait()
 
@@ -350,18 +358,13 @@ func (g *Glutton) ProduceUDP(handler string, srcAddr, dstAddr *net.UDPAddr, md c
 func (g *Glutton) Shutdown() {
 	g.cancel() // close all connection
 
-	g.Logger.Info("Shutting down listeners")
-	if err := g.Server.Shutdown(); err != nil {
-		g.Logger.Error("failed to shutdown server", producer.ErrAttr(err))
-	}
-
-	g.Logger.Info("FLushing TCP iptables")
+	g.Logger.Info("Flushing TCP iptables")
 	if err := flushTProxyIPTables(viper.GetString("interface"), g.publicAddrs[0].String(), "tcp", uint32(g.Server.tcpPort), uint32(viper.GetInt("ports.ssh"))); err != nil {
-		g.Logger.Error("failed to drop tcp iptables", producer.ErrAttr(err))
+		g.Logger.Error("Failed to drop tcp iptables", producer.ErrAttr(err))
 	}
-	g.Logger.Info("FLushing UDP iptables")
+	g.Logger.Info("Flushing UDP iptables")
 	if err := flushTProxyIPTables(viper.GetString("interface"), g.publicAddrs[0].String(), "udp", uint32(g.Server.udpPort), uint32(viper.GetInt("ports.ssh"))); err != nil {
-		g.Logger.Error("failed to drop udp iptables", producer.ErrAttr(err))
+		g.Logger.Error("Failed to drop udp iptables", producer.ErrAttr(err))
 	}
 
 	g.Logger.Info("All done")

server.go

@@ -10,7 +10,7 @@ import (
 
 type Server struct {
 	tcpListener net.Listener
-	udpListener *net.UDPConn
+	udpConn     *net.UDPConn
 	tcpPort     uint
 	udpPort     uint
 }
@@ -31,26 +31,16 @@ func (s *Server) Start() error {
 	if s.tcpListener, err = tproxy.ListenTCP("tcp4", tcpAddr); err != nil {
 		return err
 	}
+
 	udpAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("127.0.0.1:%d", s.udpPort))
 	if err != nil {
 		return err
 	}
-	if s.udpListener, err = tproxy.ListenUDP("udp4", udpAddr); err != nil {
+	if s.udpConn, err = tproxy.ListenUDP("udp4", udpAddr); err != nil {
 		return err
 	}
-	if s.udpListener == nil {
+	if s.udpConn == nil {
 		return errors.New("nil udp listener")
 	}
 	return nil
 }
-
-func (s *Server) Shutdown() error {
-	var err error
-	if s.tcpListener != nil {
-		err = s.tcpListener.Close()
-	}
-	if s.udpListener != nil {
-		err = s.udpListener.Close()
-	}
-	return err
-}

server_test.go

@@ -9,6 +9,4 @@ import (
 func TestServer(t *testing.T) {
 	server := NewServer(1234, 1235)
 	require.NotNil(t, server)
-	err := server.Shutdown()
-	require.NoError(t, err)
 }