From b3064f553c33227725a367ad7d21f1458258c7ec Mon Sep 17 00:00:00 2001
From: Jiayu Yi <yi-jiayu@users.noreply.github.com>
Date: Sat, 18 Aug 2018 00:03:05 +0800
Subject: [PATCH] Update usage

---
 README.md | 23 ++++++++++++++++++-----
 main.go   | 44 +++++++++++++++++++++++++++++---------------
 2 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/README.md b/README.md
index 569620a..6b45fe5 100644
--- a/README.md
+++ b/README.md
@@ -9,20 +9,33 @@ I wanted HTTPS for `godoc -http :6060`.
 
 ## Usage
 ```
-secure -key-file path/to/key/file -cert-file path/to/cert/file -upstream http://localhost:6060 -addr :443
+usage: secure [-addr host:port] -cert certfile -key keyfile upstream
+  -addr string
+        listen address (default ":443")
+  -cert string
+        path to cert file
+  -key string
+        path to key file
+  upstream string
+        upstream url
+```
+
+### Example
+```
+secure -cert cert.pem -key key.pem http://localhost:6060
 ```
 
 ## Demo
 *nix:
 ```
-# generate cert
+# generate self-signed certificate and private key
 openssl req -newkey rsa:4096 -nodes -keyout key.pem -x509 -days 365 -out cert.pem -subj "/CN=localhost"
 
 # start godoc
 godoc -http localhost:6060 &
 
 # secure it
-go run main.go -key-file key.pem -cert-file cert.pem -upstream http://localhost:6060 -addr :443
+secure -key key.pem -cert cert.pem http://localhost:6060
 ```
 
 Windows (PowerShell)
@@ -30,11 +43,11 @@ Windows (PowerShell)
 # somehow obtain key.pem and cert.pem
 
 # start godoc
-# cmd: start godoc -http localhost:6060
+# Command Prompt: start godoc -http localhost:6060
 Start-Process godoc "-http localhost:6060"
 
 # secure it
-go run main.go -key-file key.pem -cert-file cert.pem -upstream http://localhost:6060 -addr :443
+secure -key key.pem -cert cert.pem http://localhost:6060
 ```
 
 ## Features
diff --git a/main.go b/main.go
index 7659487..383aac1 100644
--- a/main.go
+++ b/main.go
@@ -1,14 +1,17 @@
+// secure is a super simple TLS termination proxy
 package main
 
 import (
+	"context"
 	"flag"
+	"fmt"
+	"log"
+	"net/http"
 	"net/http/httputil"
 	"net/url"
-	"fmt"
-	"net/http"
-	"os/signal"
 	"os"
-	"context"
+	"os/signal"
+	"path/filepath"
 	"syscall"
 )
 
@@ -20,15 +23,29 @@ var (
 )
 
 func init() {
-	flag.StringVar(&certFile, "cert-file", "", "path to cert file")
-	flag.StringVar(&keyFile, "key-file", "", "path to key file")
-	flag.StringVar(&upstream, "upstream", "", "upstream address")
 	flag.StringVar(&addr, "addr", ":443", "listen address")
+	flag.StringVar(&certFile, "cert", "", "path to cert file")
+	flag.StringVar(&keyFile, "key", "", "path to key file")
+
+	flag.Usage = func() {
+		fmt.Fprintf(flag.CommandLine.Output(),
+			"usage: %s [-addr host:port] -cert certfile -key keyfile upstream\n",
+			filepath.Base(os.Args[0]))
+		flag.PrintDefaults()
+		fmt.Fprintln(flag.CommandLine.Output(), "  upstream string\n    \tupstream url")
+	}
 }
 
 func _main() error {
 	flag.Parse()
 
+	if flag.NArg() == 1 {
+		upstream = flag.Arg(0)
+	} else {
+		flag.Usage()
+		os.Exit(2)
+	}
+
 	u, err := url.Parse(upstream)
 	if err != nil {
 		return fmt.Errorf("invalid upstream address: %v", err)
@@ -40,27 +57,24 @@ func _main() error {
 		Addr:    addr,
 	}
 
-	idleConnsClosed := make(chan struct{})
+	done := make(chan struct{})
 	go func() {
 		sig := make(chan os.Signal, 1)
 		signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
 		fmt.Println(<-sig)
 
-		// We received an interrupt signal, shut down.
 		if err := srv.Shutdown(context.Background()); err != nil {
-			// Error from closing listeners, or context timeout:
-			fmt.Printf("HTTP server Shutdown: %v", err)
+			fmt.Printf("Shutdown: %v", err)
 		}
-		close(idleConnsClosed)
+		close(done)
 	}()
 
+	log.Printf("cert-file=%s key-file=%s listen-addr=%s upstream-url=%s", certFile, keyFile, srv.Addr, u.String())
 	if err := srv.ListenAndServeTLS(certFile, keyFile); err != http.ErrServerClosed {
-		// Error starting or closing listener:
 		return fmt.Errorf("ListenAndServeTLS: %v", err)
 	}
 
-	<-idleConnsClosed
-
+	<-done
 	return nil
 }