Make process and czar client graceful exit

cmd/daze/main.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/mohanson/daze"
 	"github.com/mohanson/daze/lib/doa"
+	"github.com/mohanson/daze/lib/gracefulexit"
 	"github.com/mohanson/daze/protocol/ashe"
 	"github.com/mohanson/daze/protocol/baboon"
 	"github.com/mohanson/daze/protocol/czar"
@@ -111,7 +112,9 @@ func main() {
 			log.Println("main: listen net/http/pprof on", *flGpprof)
 			go func() { doa.Nil(http.ListenAndServe(*flGpprof, nil)) }()
 		}
-		daze.Hang()
+		// Hang prevent program from exiting.
+		gracefulexit.Wait()
+		log.Println("main: exit")
 	case "client":
 		var (
 			flCIDRls = flag.String("c", filepath.Join(resExec, Conf.PathCIDR), "cidr path")
@@ -160,6 +163,7 @@ func main() {
 			doa.Nil(locale.Run())
 		case "czar":
 			client := czar.NewClient(*flServer, *flCipher)
+			defer client.Close()
 			locale := daze.NewLocale(*flListen, daze.NewAimbot(client, &daze.AimbotOption{
 				Type: *flFilter,
 				Rule: *flRulels,
@@ -177,7 +181,9 @@ func main() {
 			log.Println("main: listen net/http/pprof on", *flGpprof)
 			go func() { doa.Nil(http.ListenAndServe(*flGpprof, nil)) }()
 		}
-		daze.Hang()
+		// Hang prevent program from exiting.
+		gracefulexit.Wait()
+		log.Println("main: exit")
 	case "gen":
 		flag.Usage = func() {
 			fmt.Fprint(flag.CommandLine.Output(), helpGen)

daze.go

@@ -1022,11 +1022,6 @@ func Gravity(conn io.ReadWriteCloser, k []byte) io.ReadWriteCloser {
 	}
 }
 
-// Hang prevent program from exiting.
-func Hang() {
-	select {}
-}
-
 // OpenFile select the appropriate method to open the file based on the incoming args automatically.
 //
 // Examples:
@@ -1055,22 +1050,6 @@ func (r *RandomReader) Read(p []byte) (int, error) {
 	return len(p), nil
 }
 
-// Reno is a slow start reconnection algorithm.
-func Reno(network string, address string) (net.Conn, error) {
-	i := 0
-	for {
-		r, err := Dial(network, address)
-		if err == nil {
-			return r, err
-		}
-		log.Println("reno:", err)
-		time.Sleep(time.Second * time.Duration(math.Pow(2, float64(i))))
-		if i < 5 {
-			i++
-		}
-	}
-}
-
 // Salt converts the stupid password passed in by the user to 32-sized byte array.
 func Salt(s string) []byte {
 	h := sha256.Sum256([]byte(s))

lib/gracefulexit/README.md

@@ -0,0 +1,3 @@
+# Gracefulexit
+
+A graceful exit (or graceful handling) is a simple programming idiom[citation needed] wherein a program detects a serious error condition and "exits gracefully" in a controlled manner as a result. Often the program prints a descriptive error message to a terminal or log as part of the graceful exit.

lib/gracefulexit/gracefulexit_linux.go

@@ -0,0 +1,22 @@
+// A graceful exit (or graceful handling) is a simple programming idiom[citation needed] wherein a program detects a
+// serious error condition and "exits gracefully" in a controlled manner as a result. Often the program prints a
+// descriptive error message to a terminal or log as part of the graceful exit.
+package gracefulexit
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+// Chan create a channel for os.Signal.
+func Chan() chan os.Signal {
+	buffer := make(chan os.Signal, 1)
+	signal.Notify(buffer, syscall.SIGINT, syscall.SIGTERM)
+	return buffer
+}
+
+// Wait for a signal.
+func Wait() {
+	<-Chan()
+}

lib/gracefulexit/gracefulexit_windows.go

@@ -0,0 +1,18 @@
+package gracefulexit
+
+import (
+	"os"
+	"os/signal"
+)
+
+// Chan create a channel for os.Signal.
+func Chan() chan os.Signal {
+	buffer := make(chan os.Signal, 1)
+	signal.Notify(buffer, os.Interrupt)
+	return buffer
+}
+
+// Wait for a signal.
+func Wait() {
+	<-Chan()
+}

protocol/czar/engine.go

@@ -10,7 +10,6 @@ import (
 	"time"
 
 	"github.com/mohanson/daze"
-	"github.com/mohanson/daze/lib/doa"
 	"github.com/mohanson/daze/protocol/ashe"
 )
 
@@ -114,11 +113,18 @@ func NewServer(listen string, cipher string) *Server {
 
 // Client implemented the czar protocol.
 type Client struct {
+	Cancel chan struct{}
 	Cipher []byte
 	Mux    chan *Mux
 	Server string
 }
 
+// Close the connection. All streams will be closed at the same time.
+func (c *Client) Close() error {
+	close(c.Cancel)
+	return nil
+}
+
 // Dial connects to the address on the named network.
 func (c *Client) Dial(ctx *daze.Context, network string, address string) (io.ReadWriteCloser, error) {
 	select {
@@ -140,25 +146,54 @@ func (c *Client) Dial(ctx *daze.Context, network string, address string) (io.Rea
 
 // Run creates an establish connection to czar server.
 func (c *Client) Run() {
+	var (
+		err error
+		mux *Mux
+		rtt = 0
+		sid = 0
+		srv net.Conn
+	)
 	for {
-		srv := doa.Try(daze.Reno("tcp", c.Server))
-		log.Println("czar: mux init")
-		mux := NewMuxClient(srv)
-		for {
+		switch sid {
+		case 0:
+			srv, err = daze.Dial("tcp", c.Server)
+			switch {
+			case srv == nil:
+				log.Println("czar:", err)
+				select {
+				case <-time.After(time.Second * time.Duration(math.Pow(2, float64(rtt)))):
+					// A slow start reconnection algorithm.
+					rtt = min(rtt+1, 5)
+				case <-c.Cancel:
+					sid = 2
+				}
+			case err == nil:
+				log.Println("czar: mux init")
+				mux = NewMuxClient(srv)
+				rtt = 0
+				sid = 1
+			}
+		case 1:
 			select {
 			case c.Mux <- mux:
-				continue
 			case <-mux.rdn:
+				log.Println("czar: mux done")
+				sid = 0
+			case <-c.Cancel:
+				log.Println("czar: mux done")
+				sid = 2
 			}
-			break
+		case 2:
+			mux.Close()
+			return
 		}
-		log.Println("czar: mux done")
 	}
 }
 
 // NewClient returns a new Client.
 func NewClient(server, cipher string) *Client {
 	client := &Client{
+		Cancel: make(chan struct{}),
 		Cipher: daze.Salt(cipher),
 		Mux:    make(chan *Mux),
 		Server: server,

protocol/czar/engine_test.go

@@ -24,6 +24,7 @@ func TestProtocolCzarTCP(t *testing.T) {
 	dazeServer.Run()
 
 	dazeClient := NewClient(DazeServerListenOn, Password)
+	defer dazeClient.Close()
 	ctx := &daze.Context{}
 	cli := doa.Try(dazeClient.Dial(ctx, "tcp", EchoServerListenOn))
 	defer cli.Close()
@@ -43,6 +44,7 @@ func TestProtocolCzarTCPClientClose(t *testing.T) {
 	dazeServer.Run()
 
 	dazeClient := NewClient(DazeServerListenOn, Password)
+	defer dazeClient.Close()
 	ctx := &daze.Context{}
 	cli := doa.Try(dazeClient.Dial(ctx, "tcp", EchoServerListenOn))
 	defer cli.Close()
@@ -65,6 +67,7 @@ func TestProtocolCzarTCPServerClose(t *testing.T) {
 	dazeServer.Run()
 
 	dazeClient := NewClient(DazeServerListenOn, Password)
+	defer dazeClient.Close()
 	ctx := &daze.Context{}
 	cli := doa.Try(dazeClient.Dial(ctx, "tcp", EchoServerListenOn))
 	defer cli.Close()
@@ -85,6 +88,7 @@ func TestProtocolCzarUDP(t *testing.T) {
 	dazeServer.Run()
 
 	dazeClient := NewClient(DazeServerListenOn, Password)
+	defer dazeClient.Close()
 	ctx := &daze.Context{}
 	cli := doa.Try(dazeClient.Dial(ctx, "udp", EchoServerListenOn))
 	defer cli.Close()