From 2daaa8f3ef5930d1e68a47eb7235c68211da6153 Mon Sep 17 00:00:00 2001
From: wweir <wweir@foxmail.com>
Date: Thu, 13 Dec 2018 22:27:40 +0800
Subject: [PATCH] Close session on peer going away error

---
 conf/sower.toml |  1 +
 dns/dns.go      |  4 ++--
 proxy/client.go | 12 +++++++++---
 proxy/server.go |  2 +-
 proxy/util.go   | 13 +++++++++----
 5 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/conf/sower.toml b/conf/sower.toml
index 5894acf..5364f42 100644
--- a/conf/sower.toml
+++ b/conf/sower.toml
@@ -32,6 +32,7 @@ blocklist=[
     "*.aws.amazon.com", # amazon
     "m.media-amazon.com",
     "*.awsstatic.com",
+    "*.s3.amazonaws.com",
 	"*.cloudfront.net", # atlassian
     "synchrony-cdn.atlassian.com",
     "avatar-cdn.atlassian.com",
diff --git a/dns/dns.go b/dns/dns.go
index 8c5ce07..1854a0b 100644
--- a/dns/dns.go
+++ b/dns/dns.go
@@ -76,8 +76,8 @@ func manual(w dns.ResponseWriter, r *dns.Msg, domain, dnsServer string) {
 		return
 	}
 
-	msg, err := dns.Exchange(r, dnsServer+":53") // expose any response
-	if msg == nil {
+	msg, err := dns.Exchange(r, dnsServer+":53")
+	if msg == nil { // expose any response except nil
 		glog.V(1).Infof("get dns of %s fail: %s", domain, err)
 		return
 	}
diff --git a/proxy/client.go b/proxy/client.go
index b00d6c7..e80a41a 100644
--- a/proxy/client.go
+++ b/proxy/client.go
@@ -13,6 +13,7 @@ func StartClient(server string) {
 	connCh := listenLocal([]string{":80", ":443"})
 	reDialCh := make(chan net.Conn, 10)
 	var conn net.Conn
+	var count int
 
 	for {
 		sess, err := quic.DialAddr(server, &tls.Config{InsecureSkipVerify: true}, dialConf)
@@ -25,14 +26,16 @@ func StartClient(server string) {
 		}
 		glog.Infof("new session from (%s) to (%s)", sess.LocalAddr(), sess.RemoteAddr())
 
+		count = 0
 		for { // session rotate logic
 			select {
 			case conn = <-connCh:
 			case conn = <-reDialCh:
 			}
+			count++
 
 			// sync action to reuse sigle sess
-			if !openStream(conn, sess, reDialCh) {
+			if !openStream(conn, sess, count, reDialCh) {
 				sess.Close()
 				break
 			}
@@ -40,7 +43,7 @@ func StartClient(server string) {
 	}
 }
 
-func openStream(conn net.Conn, sess quic.Session, reDialCh chan<- net.Conn) bool {
+func openStream(conn net.Conn, sess quic.Session, count int, reDialCh chan<- net.Conn) bool {
 	glog.V(2).Infoln("new request from", conn.RemoteAddr())
 
 	okCh := make(chan struct{})
@@ -54,6 +57,9 @@ func openStream(conn net.Conn, sess quic.Session, reDialCh chan<- net.Conn) bool
 		}
 		defer stream.Close()
 
+		glog.V(2).Infof("START stream\t%d", count)
+		defer glog.V(2).Infof("CLOSE stream\t%d", count)
+
 		select {
 		case okCh <- struct{}{}:
 		default:
@@ -65,7 +71,7 @@ func openStream(conn net.Conn, sess quic.Session, reDialCh chan<- net.Conn) bool
 		if err := conn.(*net.TCPConn).SetKeepAlive(true); err != nil {
 			glog.Warningln(err)
 		}
-		relay(&streamConn{stream, sess}, conn)
+		relay(sess, &streamConn{stream, sess}, conn)
 		conn.Close()
 	}()
 
diff --git a/proxy/server.go b/proxy/server.go
index 8989b4c..bb13861 100644
--- a/proxy/server.go
+++ b/proxy/server.go
@@ -59,5 +59,5 @@ func acceptStream(stream quic.Stream, sess quic.Session) {
 	if err := rc.(*net.TCPConn).SetKeepAlive(true); err != nil {
 		glog.Warningln(err)
 	}
-	relay(rc, conn)
+	relay(sess, rc, conn)
 }
diff --git a/proxy/util.go b/proxy/util.go
index 1aee6dc..ebb783b 100644
--- a/proxy/util.go
+++ b/proxy/util.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"math/big"
 	"net"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -36,18 +37,22 @@ func (s *streamConn) RemoteAddr() net.Addr {
 	return s.sess.RemoteAddr()
 }
 
-func relay(conn1, conn2 net.Conn) {
+func relay(sess quic.Session, conn1, conn2 net.Conn) {
 	wg := &sync.WaitGroup{}
 	exitFlag := new(int32)
 	wg.Add(2)
-	go redirect(conn1, conn2, wg, exitFlag)
-	redirect(conn2, conn1, wg, exitFlag)
+	go redirect(sess, conn1, conn2, wg, exitFlag)
+	redirect(sess, conn2, conn1, wg, exitFlag)
 	wg.Wait()
 }
 
-func redirect(conn1, conn2 net.Conn, wg *sync.WaitGroup, exitFlag *int32) {
+func redirect(sess quic.Session, conn1, conn2 net.Conn, wg *sync.WaitGroup, exitFlag *int32) {
 	if _, err := io.Copy(conn2, conn1); err != nil && (atomic.LoadInt32(exitFlag) == 0) {
 		glog.V(1).Infof("%s<>%s -> %s<>%s: %s", conn1.RemoteAddr(), conn1.LocalAddr(), conn2.LocalAddr(), conn2.RemoteAddr(), err)
+
+		if strings.Contains(err.Error(), "PeerGoingAway") { //for internal package, hard code here
+			sess.Close()
+		}
 	}
 
 	// wakeup all conn goroutine