add close server function

Author: ilanyu

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
+}

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())
+}

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
+		}
 	}
 }