From 791e1e22cb79344f0d0b9890a5e8133ec5d85c0d Mon Sep 17 00:00:00 2001
From: mattn <mattn.jp@gmail.com>
Date: Thu, 20 Aug 2015 15:40:38 +0900
Subject: [PATCH] HTTP Basic Authentication support. Close #8

---
 app/app.go | 42 ++++++++++++++++++++++++++++++++++++++++--
 main.go    |  7 ++++++-
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/app/app.go b/app/app.go
index ec2d321..178853f 100644
--- a/app/app.go
+++ b/app/app.go
@@ -1,6 +1,7 @@
 package app
 
 import (
+	"encoding/base64"
 	"encoding/json"
 	"log"
 	"net/http"
@@ -19,18 +20,50 @@ type App struct {
 	Address     string
 	Port        string
 	PermitWrite bool
+	Credential  string
 	Command     []string
 }
 
-func New(address string, port string, permitWrite bool, command []string) *App {
+func New(address string, port string, permitWrite bool, cred string, command []string) *App {
 	return &App{
 		Address:     address,
 		Port:        port,
 		PermitWrite: permitWrite,
+		Credential:  cred,
 		Command:     command,
 	}
 }
 
+func loggerHandler(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		log.Printf("%s %s", r.Method, r.URL.Path)
+		h.ServeHTTP(w, r)
+	})
+}
+
+func basicAuthHandler(h http.Handler, cred string) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		token := strings.SplitN(r.Header.Get("Authorization"), " ", 2)
+		if len(token) != 2 || strings.ToLower(token[0]) != "basic" {
+			w.Header().Set("WWW-Authenticate", `Basic realm="GoTTY"`)
+			http.Error(w, "Bad Request", http.StatusUnauthorized)
+			return
+		}
+
+		payload, err := base64.StdEncoding.DecodeString(token[1])
+		if err != nil {
+			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+			return
+		}
+		if cred != string(payload) {
+			w.Header().Set("WWW-Authenticate", `Basic realm="GoTTY"`)
+			http.Error(w, "authorization failed", http.StatusUnauthorized)
+			return
+		}
+		h.ServeHTTP(w, r)
+	})
+}
+
 func (app *App) Run() error {
 	http.Handle("/",
 		http.FileServer(
@@ -41,7 +74,12 @@ func (app *App) Run() error {
 
 	url := app.Address + ":" + app.Port
 	log.Printf("Server is running at %s, command: %s", url, strings.Join(app.Command, " "))
-	err := http.ListenAndServe(url, nil)
+	handler := http.Handler(http.DefaultServeMux)
+	handler = loggerHandler(handler)
+	if app.Credential != "" {
+		handler = basicAuthHandler(handler, app.Credential)
+	}
+	err := http.ListenAndServe(url, handler)
 	if err != nil {
 		return err
 	}
diff --git a/main.go b/main.go
index f689c36..1f8fc63 100644
--- a/main.go
+++ b/main.go
@@ -32,6 +32,11 @@ func main() {
 			Usage:  "Permit clients to write to the TTY (BE CAREFUL)",
 			EnvVar: "GOTTY_PERMIT_WRITE",
 		},
+		cli.StringFlag{
+			Name:   "credential, c",
+			Usage:  "Credential for Basic Authentication (ex: user:pass)",
+			EnvVar: "GOTTY_CREDENTIAL",
+		},
 	}
 	cmd.Action = func(c *cli.Context) {
 		if len(c.Args()) == 0 {
@@ -39,7 +44,7 @@ func main() {
 			cli.ShowAppHelp(c)
 			os.Exit(1)
 		}
-		app := app.New(c.String("addr"), c.String("port"), c.Bool("permit-write"), c.Args())
+		app := app.New(c.String("addr"), c.String("port"), c.Bool("permit-write"), c.String("credential"), c.Args())
 		err := app.Run()
 		if err != nil {
 			fmt.Println(err)