academy-chauffeurs - auth.go

     1  package main
     2  
     3  import (
     4  	"crypto/sha1"
     5  	"net/http"
     6  	"sync"
     7  
     8  	"golang.org/x/net/websocket"
     9  )
    10  
    11  type noCache struct {
    12  	http.Handler
    13  }
    14  
    15  func (n noCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    16  	if r.URL.Path == "/" {
    17  		w.Header().Add("Cache-Control", "no-cache, no-store, must-revalidate")
    18  	}
    19  	n.Handler.ServeHTTP(w, r)
    20  }
    21  
    22  type UserMap struct {
    23  	sync.RWMutex
    24  	users map[string]uint
    25  }
    26  
    27  var userMap UserMap
    28  
    29  func (u *UserMap) Add(username string) {
    30  	u.Lock()
    31  	u.users[username] = u.users[username] + 1
    32  	u.Unlock()
    33  }
    34  
    35  func (u *UserMap) Remove(username string) {
    36  	u.Lock()
    37  	count := u.users[username] - 1
    38  	if count == 0 {
    39  		delete(u.users, username)
    40  	} else {
    41  		u.users[username] = count
    42  	}
    43  	u.Unlock()
    44  }
    45  
    46  func (u *UserMap) Copy() map[string]uint {
    47  	u.RLock()
    48  	m := make(map[string]uint, len(u.users))
    49  	for u, c := range u.users {
    50  		m[u] = c
    51  	}
    52  	u.RUnlock()
    53  	return m
    54  }
    55  
    56  type userConn struct {
    57  	websocket.Handler
    58  }
    59  
    60  func (u *userConn) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    61  	user, _, _ := r.BasicAuth()
    62  	userMap.Add(user)
    63  	u.Handler.ServeHTTP(w, r)
    64  	userMap.Remove(user)
    65  }
    66  
    67  type AuthMap struct {
    68  	sync.RWMutex
    69  	users map[string][sha1.Size]byte
    70  }
    71  
    72  var authMap AuthMap
    73  
    74  func init() {
    75  	userMap.users = make(map[string]uint)
    76  	authMap.users = make(map[string][sha1.Size]byte)
    77  }
    78  
    79  // Set sets the password for the given username. Returns true if the user
    80  // already exits.
    81  func (a *AuthMap) Set(username string, password [sha1.Size]byte) bool {
    82  	a.Lock()
    83  	_, ok := a.users[username]
    84  	a.users[username] = password
    85  	a.Unlock()
    86  	return ok
    87  }
    88  
    89  func (a *AuthMap) Check(username, password string) bool {
    90  	pwd := sha1.Sum([]byte(password))
    91  	a.RLock()
    92  	p, ok := a.users[username]
    93  	a.RUnlock()
    94  	return ok && pwd == p
    95  }
    96  
    97  func (a *AuthMap) Remove(username string) {
    98  	a.Lock()
    99  	delete(a.users, username)
   100  	a.Unlock()
   101  }
   102  
   103  type authServeMux struct {
   104  	http.ServeMux
   105  }
   106  
   107  func (a *authServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
   108  	username, password, ok := r.BasicAuth()
   109  	if ok {
   110  		ok = authMap.Check(username, password)
   111  	}
   112  	if !ok {
   113  		w.Header().Set("WWW-Authenticate", "Basic realm=\"Enter Credentials\"")
   114  		w.WriteHeader(http.StatusUnauthorized)
   115  		w.Write(unauthorised)
   116  		return
   117  	}
   118  	a.ServeMux.ServeHTTP(w, r)
   119  }
   120  
   121  var unauthorised = []byte(`<html>
   122  	<head>
   123  		<title>Unauthorised</title>
   124  	</head>
   125  	<body>
   126  		<h1>Not Authorised</h1>
   127  	</body>
   128  </html>
   129  `)