From c3d1b02219dae6999c78dad5bf1aec21b5d72821 Mon Sep 17 00:00:00 2001
From: venjiang <venjiang@gmail.com>
Date: Wed, 3 Aug 2022 17:07:13 +0800
Subject: [PATCH] feat: add server connection handler

---
 core/server.go | 37 ++++++++++++++++++++++++++-----------
 1 file changed, 26 insertions(+), 11 deletions(-)

diff --git a/core/server.go b/core/server.go
index 7e04de9ce..2920b002a 100644
--- a/core/server.go
+++ b/core/server.go
@@ -31,19 +31,23 @@ type ServerOption func(*ServerOptions)
 // FrameHandler is the handler for frame.
 type FrameHandler func(c *Context) error
 
+// ConnectionHandler is the handler for quic connection
+type ConnectionHandler func(conn quic.Connection)
+
 // Server is the underlining server of Zipper
 type Server struct {
-	name               string
-	state              string
-	connector          Connector
-	router             Router
-	metadataBuilder    MetadataBuilder
-	counterOfDataFrame int64
-	downstreams        map[string]*Client
-	mu                 sync.Mutex
-	opts               ServerOptions
-	beforeHandlers     []FrameHandler
-	afterHandlers      []FrameHandler
+	name                    string
+	state                   string
+	connector               Connector
+	router                  Router
+	metadataBuilder         MetadataBuilder
+	counterOfDataFrame      int64
+	downstreams             map[string]*Client
+	mu                      sync.Mutex
+	opts                    ServerOptions
+	beforeHandlers          []FrameHandler
+	afterHandlers           []FrameHandler
+	connectionCloseHandlers []ConnectionHandler
 }
 
 // NewServer create a Server instance.
@@ -120,6 +124,12 @@ func (s *Server) Serve(ctx context.Context, conn net.PacketConn) error {
 		logger.Infof("%s❤️1/ new connection: %s", ServerLogPrefix, connID)
 
 		go func(ctx context.Context, qconn quic.Connection) {
+			// connection close handlers
+			defer func() {
+				for _, h := range s.connectionCloseHandlers {
+					h(qconn)
+				}
+			}()
 			for {
 				logger.Infof("%s❤️2/ waiting for new stream", ServerLogPrefix)
 				stream, err := qconn.AcceptStream(ctx)
@@ -532,6 +542,11 @@ func (s *Server) SetAfterHandlers(handlers ...FrameHandler) {
 	s.afterHandlers = append(s.afterHandlers, handlers...)
 }
 
+// SetConnectionCloseHandlers set the connection close handlers of server.
+func (s *Server) SetConnectionCloseHandlers(handlers ...ConnectionHandler) {
+	s.connectionCloseHandlers = append(s.connectionCloseHandlers, handlers...)
+}
+
 func (s *Server) authNames() []string {
 	if len(s.opts.Auths) == 0 {
 		return []string{"none"}