services/keep/src/arvados.org/keepproxy/keepproxy.go

   1 package main
   2
   3 import (
   4         "arvados.org/keepclient"
   5         "flag"
   6         "fmt"
   7         "github.com/gorilla/mux"
   8         "io"
   9         "log"
  10         "net"
  11         "net/http"
  12         "os"
  13         "sync"
  14         "time"
  15 )
  16
  17 // Default TCP address on which to listen for requests.
  18 // Initialized by the -listen flag.
  19 const DEFAULT_ADDR = ":25107"
  20
  21 var listener net.Listener
  22
  23 func main() {
  24         var (
  25                 listen           string
  26                 no_get           bool
  27                 no_put           bool
  28                 default_replicas int
  29                 pidfile          string
  30         )
  31
  32         flag.StringVar(
  33                 &listen,
  34                 "listen",
  35                 DEFAULT_ADDR,
  36                 "Interface on which to listen for requests, in the format "+
  37                         "ipaddr:port. e.g. -listen=10.0.1.24:8000. Use -listen=:port "+
  38                         "to listen on all network interfaces.")
  39
  40         flag.BoolVar(
  41                 &no_get,
  42                 "no-get",
  43                 false,
  44                 "If set, disable GET operations")
  45
  46         flag.BoolVar(
  47                 &no_get,
  48                 "no-put",
  49                 false,
  50                 "If set, disable PUT operations")
  51
  52         flag.IntVar(
  53                 &default_replicas,
  54                 "default-replicas",
  55                 2,
  56                 "Default number of replicas to write if not specified by the client.")
  57
  58         flag.StringVar(
  59                 &pidfile,
  60                 "pid",
  61                 "",
  62                 "Path to write pid file")
  63
  64         flag.Parse()
  65
  66         kc, err := keepclient.MakeKeepClient()
  67         if err != nil {
  68                 log.Fatal(err)
  69         }
  70
  71         if pidfile != "" {
  72                 f, err := os.Create(pidfile)
  73                 if err == nil {
  74                         fmt.Fprint(f, os.Getpid())
  75                         f.Close()
  76                 } else {
  77                         log.Printf("Error writing pid file (%s): %s", pidfile, err.Error())
  78                 }
  79         }
  80
  81         kc.Want_replicas = default_replicas
  82
  83         listener, err = net.Listen("tcp", listen)
  84         if err != nil {
  85                 log.Fatalf("Could not listen on %v", listen)
  86         }
  87
  88         go RefreshServicesList(&kc)
  89
  90         // Start listening for requests.
  91         http.Serve(listener, MakeRESTRouter(!no_get, !no_put, &kc))
  92 }
  93
  94 type ApiTokenCache struct {
  95         tokens     map[string]int64
  96         lock       sync.Mutex
  97         expireTime int64
  98 }
  99
 100 // Refresh the keep service list every five minutes.
 101 func RefreshServicesList(kc *keepclient.KeepClient) {
 102         for {
 103                 time.Sleep(300 * time.Second)
 104                 kc.DiscoverKeepServers()
 105         }
 106 }
 107
 108 // Cache the token and set an expire time.  If we already have an expire time
 109 // on the token, it is not updated.
 110 func (this *ApiTokenCache) RememberToken(token string) {
 111         this.lock.Lock()
 112         defer this.lock.Unlock()
 113
 114         now := time.Now().Unix()
 115         if this.tokens[token] == 0 {
 116                 this.tokens[token] = now + this.expireTime
 117         }
 118 }
 119
 120 // Check if the cached token is known and still believed to be valid.
 121 func (this *ApiTokenCache) RecallToken(token string) bool {
 122         this.lock.Lock()
 123         defer this.lock.Unlock()
 124
 125         now := time.Now().Unix()
 126         if this.tokens[token] == 0 {
 127                 // Unknown token
 128                 return false
 129         } else if now < this.tokens[token] {
 130                 // Token is known and still valid
 131                 return true
 132         } else {
 133                 // Token is expired
 134                 this.tokens[token] = 0
 135                 return false
 136         }
 137 }
 138
 139 func GetRemoteAddress(req *http.Request) string {
 140         if realip := req.Header.Get("X-Real-IP"); realip != "" {
 141                 if forwarded := req.Header.Get("X-Forwarded-For"); forwarded != realip {
 142                         return fmt.Sprintf("%s (X-Forwarded-For %s)", realip, forwarded)
 143                 } else {
 144                         return realip
 145                 }
 146         }
 147         return req.RemoteAddr
 148 }
 149
 150 func CheckAuthorizationHeader(kc keepclient.KeepClient, cache *ApiTokenCache, req *http.Request) bool {
 151         var auth string
 152         if auth = req.Header.Get("Authorization"); auth == "" {
 153                 return false
 154         }
 155
 156         var tok string
 157         _, err := fmt.Sscanf(auth, "OAuth2 %s", &tok)
 158         if err != nil {
 159                 // Scanning error
 160                 return false
 161         }
 162
 163         if cache.RecallToken(tok) {
 164                 // Valid in the cache, short circut
 165                 return true
 166         }
 167
 168         var usersreq *http.Request
 169
 170         if usersreq, err = http.NewRequest("GET", fmt.Sprintf("https://%s/arvados/v1/users/current", kc.ApiServer), nil); err != nil {
 171                 // Can't construct the request
 172                 log.Printf("%s: CheckAuthorizationHeader error: %v", GetRemoteAddress(req), err)
 173                 return false
 174         }
 175
 176         // Add api token header
 177         usersreq.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", tok))
 178
 179         // Actually make the request
 180         var resp *http.Response
 181         if resp, err = kc.Client.Do(usersreq); err != nil {
 182                 // Something else failed
 183                 log.Printf("%s: CheckAuthorizationHeader error connecting to API server: %v", GetRemoteAddress(req), err.Error())
 184                 return false
 185         }
 186
 187         if resp.StatusCode != http.StatusOK {
 188                 // Bad status
 189                 log.Printf("%s: CheckAuthorizationHeader API server responded: %v", GetRemoteAddress(req), resp.Status)
 190                 return false
 191         }
 192
 193         // Success!  Update cache
 194         cache.RememberToken(tok)
 195
 196         return true
 197 }
 198
 199 type GetBlockHandler struct {
 200         *keepclient.KeepClient
 201         *ApiTokenCache
 202 }
 203
 204 type PutBlockHandler struct {
 205         *keepclient.KeepClient
 206         *ApiTokenCache
 207 }
 208
 209 // MakeRESTRouter
 210 //     Returns a mux.Router that passes GET and PUT requests to the
 211 //     appropriate handlers.
 212 //
 213 func MakeRESTRouter(
 214         enable_get bool,
 215         enable_put bool,
 216         kc *keepclient.KeepClient) *mux.Router {
 217
 218         t := &ApiTokenCache{tokens: make(map[string]int64), expireTime: 300}
 219
 220         rest := mux.NewRouter()
 221         gh := rest.Handle(`/{hash:[0-9a-f]{32}}`, GetBlockHandler{kc, t})
 222         ghsig := rest.Handle(
 223                 `/{hash:[0-9a-f]{32}}+A{signature:[0-9a-f]+}@{timestamp:[0-9a-f]+}`,
 224                 GetBlockHandler{kc, t})
 225         ph := rest.Handle(`/{hash:[0-9a-f]{32}}`, PutBlockHandler{kc, t})
 226
 227         if enable_get {
 228                 gh.Methods("GET", "HEAD")
 229                 ghsig.Methods("GET", "HEAD")
 230         }
 231
 232         if enable_put {
 233                 ph.Methods("PUT")
 234         }
 235
 236         return rest
 237 }
 238
 239 func (this GetBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 240
 241         kc := *this.KeepClient
 242
 243         hash := mux.Vars(req)["hash"]
 244         signature := mux.Vars(req)["signature"]
 245         timestamp := mux.Vars(req)["timestamp"]
 246
 247         log.Printf("%s: %s %s", GetRemoteAddress(req), req.Method, hash)
 248
 249         if !CheckAuthorizationHeader(kc, this.ApiTokenCache, req) {
 250                 http.Error(resp, "Missing or invalid Authorization header", http.StatusForbidden)
 251                 return
 252         }
 253
 254         var reader io.ReadCloser
 255         var err error
 256         var blocklen int64
 257
 258         if req.Method == "GET" {
 259                 reader, blocklen, _, err = kc.AuthorizedGet(hash, signature, timestamp)
 260                 defer reader.Close()
 261         } else if req.Method == "HEAD" {
 262                 blocklen, _, err = kc.AuthorizedAsk(hash, signature, timestamp)
 263         }
 264
 265         resp.Header().Set("Content-Length", fmt.Sprint(blocklen))
 266
 267         switch err {
 268         case nil:
 269                 if reader != nil {
 270                         n, err2 := io.Copy(resp, reader)
 271                         if n != blocklen {
 272                                 log.Printf("%s: %s %s mismatched return %v with Content-Length %v error", GetRemoteAddress(req), req.Method, hash, n, blocklen, err.Error())
 273                         } else if err2 == nil {
 274                                 log.Printf("%s: %s %s success returned %v bytes", GetRemoteAddress(req), req.Method, hash, n)
 275                         } else {
 276                                 log.Printf("%s: %s %s returned %v bytes error %v", GetRemoteAddress(req), req.Method, hash, n, err.Error())
 277                         }
 278                 } else {
 279                         log.Printf("%s: %s %s success", GetRemoteAddress(req), req.Method, hash)
 280                 }
 281         case keepclient.BlockNotFound:
 282                 http.Error(resp, "Not found", http.StatusNotFound)
 283         default:
 284                 http.Error(resp, err.Error(), http.StatusBadGateway)
 285         }
 286
 287         if err != nil {
 288                 log.Printf("%s: %s %s error %s", GetRemoteAddress(req), req.Method, hash, err.Error())
 289         }
 290 }
 291
 292 func (this PutBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 293
 294         kc := *this.KeepClient
 295
 296         hash := mux.Vars(req)["hash"]
 297
 298         var contentLength int64 = -1
 299         if req.Header.Get("Content-Length") != "" {
 300                 _, err := fmt.Sscanf(req.Header.Get("Content-Length"), "%d", &contentLength)
 301                 if err != nil {
 302                         resp.Header().Set("Content-Length", fmt.Sprintf("%d", contentLength))
 303                 }
 304
 305         }
 306
 307         log.Printf("%s: %s %s Content-Length %v", GetRemoteAddress(req), req.Method, hash, contentLength)
 308
 309         if contentLength < 1 {
 310                 http.Error(resp, "Must include Content-Length header", http.StatusLengthRequired)
 311                 return
 312         }
 313
 314         if !CheckAuthorizationHeader(kc, this.ApiTokenCache, req) {
 315                 http.Error(resp, "Missing or invalid Authorization header", http.StatusForbidden)
 316                 return
 317         }
 318
 319         // Check if the client specified the number of replicas
 320         if req.Header.Get("X-Keep-Desired-Replicas") != "" {
 321                 var r int
 322                 _, err := fmt.Sscanf(req.Header.Get(keepclient.X_Keep_Desired_Replicas), "%d", &r)
 323                 if err != nil {
 324                         kc.Want_replicas = r
 325                 }
 326         }
 327
 328         // Now try to put the block through
 329         replicas, err := kc.PutHR(hash, req.Body, contentLength)
 330
 331         // Tell the client how many successful PUTs we accomplished
 332         resp.Header().Set(keepclient.X_Keep_Replicas_Stored, fmt.Sprintf("%d", replicas))
 333
 334         switch err {
 335         case nil:
 336                 // Default will return http.StatusOK
 337                 log.Printf("%s: %s %s finished, stored %v replicas (desired %v)", GetRemoteAddress(req), req.Method, hash, replicas, kc.Want_replicas)
 338
 339         case keepclient.OversizeBlockError:
 340                 // Too much data
 341                 http.Error(resp, fmt.Sprintf("Exceeded maximum blocksize %d", keepclient.BLOCKSIZE), http.StatusRequestEntityTooLarge)
 342
 343         case keepclient.InsufficientReplicasError:
 344                 if replicas > 0 {
 345                         // At least one write is considered success.  The
 346                         // client can decide if getting less than the number of
 347                         // replications it asked for is a fatal error.
 348                         // Default will return http.StatusOK
 349                 } else {
 350                         http.Error(resp, "", http.StatusServiceUnavailable)
 351                 }
 352
 353         default:
 354                 http.Error(resp, err.Error(), http.StatusBadGateway)
 355         }
 356
 357         if err != nil {
 358                 log.Printf("%s: %s %s stored %v replicas (desired %v) got error %v", GetRemoteAddress(req), req.Method, hash, replicas, kc.Want_replicas, err.Error())
 359         }
 360
 361 }