From e2fc89f1af0067704c53ed9acb3ff03c84a8892b Mon Sep 17 00:00:00 2001
From: Masahiro Nagano <kazeburo@gmail.com>
Date: Thu, 11 Oct 2018 08:54:05 +0900
Subject: [PATCH] log about disconnect

---
 wsgate-server.go | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/wsgate-server.go b/wsgate-server.go
index fce43f8..9e6bbc6 100644
--- 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)