log about disconnect

kazeburo · Oct 10, 2018 · e2fc89f · e2fc89f
1 parent f2b48cb
commit e2fc89f
Showing 1 changed file with 16 additions and 2 deletions.
diff --git a/wsgate-server.go b/wsgate-server.go
@@ -49,6 +49,7 @@ func handleProxy(w http.ResponseWriter, r *http.Request) {
 	readLen := int64(0)
 	writeLen := int64(0)
 	hasError := false
+	disconnectAt := ""
 
 	if *publicKeyFile != "" {
 		tokenString := r.Header.Get("Authorization")
@@ -84,6 +85,7 @@ func handleProxy(w http.ResponseWriter, r *http.Request) {
 	}
 
 	s, err := net.DialTimeout("tcp", upstream, *dialTimeout)
+
 	if err != nil {
 		hasError = true
 		log.Printf("DialTimeout: %v", err)
@@ -107,9 +109,9 @@ func handleProxy(w http.ResponseWriter, r *http.Request) {
 		if hasError {
 			status = "Failed"
 		}
-		log.Printf("status:%s dest:%s upstream:%s x-forwarded-for:%s remote_addr:%s read:%d write:%d",
+		log.Printf("status:%s dest:%s upstream:%s x-forwarded-for:%s remote_addr:%s read:%d write:%d disconnect_at:%s",
 			status, proxyDest, upstream, r.Header.Get("X-Forwarded-For"),
-			r.RemoteAddr, readLen, writeLen)
+			r.RemoteAddr, readLen, writeLen, disconnectAt)
 	}()
 
 	doneCh := make(chan bool)
@@ -131,6 +133,9 @@ func handleProxy(w http.ResponseWriter, r *http.Request) {
 					log.Printf("NextReader: %v", err)
 					hasError = true
 				}
+				if disconnectAt == "" {
+					disconnectAt = "client_nextreader"
+				}
 				return
 			}
 			if mt != websocket.BinaryMessage {
@@ -144,6 +149,9 @@ func handleProxy(w http.ResponseWriter, r *http.Request) {
 					log.Printf("Reading from websocket: %v", err)
 					hasError = true
 				}
+				if disconnectAt == "" {
+					disconnectAt = "client_upstream_copy"
+				}
 				return
 			}
 			readLen += n
@@ -161,6 +169,9 @@ func handleProxy(w http.ResponseWriter, r *http.Request) {
 					log.Printf("Reading from dest: %v", err)
 					hasError = true
 				}
+				if disconnectAt == "" {
+					disconnectAt = "upstream_read"
+				}
 				return
 			}
 
@@ -171,6 +182,9 @@ func handleProxy(w http.ResponseWriter, r *http.Request) {
 					log.Printf("WriteMessage: %v", err)
 					hasError = true
 				}
+				if disconnectAt == "" {
+					disconnectAt = "client_write"
+				}
 				return
 			}
 			writeLen += int64(n)