From 345df7cd9821754bad30c008c0a696b42a1f1aaf Mon Sep 17 00:00:00 2001
From: ilanyu <lanyu19950316@gmail.com>
Date: Fri, 23 Mar 2018 14:23:34 +0800
Subject: [PATCH] add close server function

---
 cmd.go    | 16 ++++++++++++++++
 handle.go | 23 +++++++++++++++++++++++
 main.go   | 52 +++++++++++++++++++++++++++++++---------------------
 3 files changed, 70 insertions(+), 21 deletions(-)
 create mode 100644 cmd.go
 create mode 100644 handle.go

diff --git a/cmd.go b/cmd.go
new file mode 100644
index 0000000..125673f
--- /dev/null
+++ b/cmd.go
@@ -0,0 +1,16 @@
+package main
+
+import "flag"
+
+type Cmd struct {
+	bind string
+	remote string
+}
+
+func parseCmd() Cmd {
+	var cmd Cmd
+	flag.StringVar(&cmd.bind, "l", "0.0.0.0:8888", "listen on ip:port")
+	flag.StringVar(&cmd.remote, "r", "http://idea.lanyus.com:80", "reverse proxy addr")
+	flag.Parse()
+	return cmd
+}
\ No newline at end of file
diff --git a/handle.go b/handle.go
new file mode 100644
index 0000000..88560de
--- /dev/null
+++ b/handle.go
@@ -0,0 +1,23 @@
+package main
+
+import (
+	"net/http"
+	"net/url"
+	"net/http/httputil"
+	"log"
+)
+
+type handle struct {
+	reverseProxy string
+}
+
+func (this *handle) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	remote, err := url.Parse(this.reverseProxy)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	proxy := httputil.NewSingleHostReverseProxy(remote)
+	r.Host = remote.Host
+	proxy.ServeHTTP(w, r)
+	log.Println(r.RemoteAddr + " " + r.Method + " " + r.URL.String() + " " + r.Proto + " " + r.UserAgent())
+}
diff --git a/main.go b/main.go
index 5cd0125..b9dc47e 100644
--- a/main.go
+++ b/main.go
@@ -3,34 +3,44 @@ package main
 import (
 	"log"
 	"net/http"
-	"net/http/httputil"
-	"net/url"
-	"flag"
+	"bufio"
+	"os"
+	"strings"
 )
 
-type handle struct {
-	reverseProxy string
+var srv http.Server
+
+func StartServer(bind string, remote string)  {
+	log.Printf("Listening on %s, forwarding to %s", bind, remote)
+	h := &handle{reverseProxy: remote}
+	srv.Addr = bind
+	srv.Handler = h
+	go func() {
+		if err := srv.ListenAndServe(); err != nil {
+			log.Fatalln("ListenAndServe: ", err)
+		}
+	}()
 }
 
-func (this *handle) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	remote, err := url.Parse(this.reverseProxy)
-	if err != nil {
-		log.Fatalln(err)
+func StopServer()  {
+	if err := srv.Close() ; err != nil {
+		log.Println(err)
 	}
-	proxy := httputil.NewSingleHostReverseProxy(remote)
-	r.Host = remote.Host
-	proxy.ServeHTTP(w, r)
-	log.Println(r.RemoteAddr + " " + r.Method + " " + r.URL.String() + " " + r.Proto + " " + r.UserAgent())
 }
 
 func main() {
-	bind := flag.String("l", "0.0.0.0:8888", "listen on ip:port")
-	remote := flag.String("r", "http://idea.lanyus.com:80", "reverse proxy addr")
-	flag.Parse()
-	log.Printf("Listening on %s, forwarding to %s", *bind, *remote)
-	h := &handle{reverseProxy: *remote}
-	err := http.ListenAndServe(*bind, h)
-	if err != nil {
-		log.Fatalln("ListenAndServe: ", err)
+	cmd := parseCmd()
+	StartServer(cmd.bind, cmd.remote)
+	reader := bufio.NewReader(os.Stdin)
+	for {
+		str , err := reader.ReadString('\n')
+		if err != nil {
+			log.Println(err)
+		}
+		if strings.TrimSpace(str) == "stop" {
+			log.Println("will stop server")
+			StopServer()
+			return
+		}
 	}
 }