services/keepproxy/keepproxy.go

   1 package main
   2
   3 import (
   4         "errors"
   5         "flag"
   6         "fmt"
   7         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
   8         "git.curoverse.com/arvados.git/sdk/go/keepclient"
   9         "github.com/gorilla/mux"
  10         "io"
  11         "io/ioutil"
  12         "log"
  13         "net"
  14         "net/http"
  15         "os"
  16         "os/signal"
  17         "sync"
  18         "syscall"
  19         "time"
  20 )
  21
  22 // Default TCP address on which to listen for requests.
  23 // Initialized by the -listen flag.
  24 const DEFAULT_ADDR = ":25107"
  25
  26 var listener net.Listener
  27
  28 func main() {
  29         var (
  30                 listen           string
  31                 no_get           bool
  32                 no_put           bool
  33                 default_replicas int
  34                 timeout          int64
  35                 pidfile          string
  36         )
  37
  38         flagset := flag.NewFlagSet("default", flag.ExitOnError)
  39
  40         flagset.StringVar(
  41                 &listen,
  42                 "listen",
  43                 DEFAULT_ADDR,
  44                 "Interface on which to listen for requests, in the format "+
  45                         "ipaddr:port. e.g. -listen=10.0.1.24:8000. Use -listen=:port "+
  46                         "to listen on all network interfaces.")
  47
  48         flagset.BoolVar(
  49                 &no_get,
  50                 "no-get",
  51                 false,
  52                 "If set, disable GET operations")
  53
  54         flagset.BoolVar(
  55                 &no_put,
  56                 "no-put",
  57                 false,
  58                 "If set, disable PUT operations")
  59
  60         flagset.IntVar(
  61                 &default_replicas,
  62                 "default-replicas",
  63                 2,
  64                 "Default number of replicas to write if not specified by the client.")
  65
  66         flagset.Int64Var(
  67                 &timeout,
  68                 "timeout",
  69                 15,
  70                 "Timeout on requests to internal Keep services (default 15 seconds)")
  71
  72         flagset.StringVar(
  73                 &pidfile,
  74                 "pid",
  75                 "",
  76                 "Path to write pid file")
  77
  78         flagset.Parse(os.Args[1:])
  79
  80         arv, err := arvadosclient.MakeArvadosClient()
  81         if err != nil {
  82                 log.Fatalf("Error setting up arvados client %s", err.Error())
  83         }
  84
  85         kc, err := keepclient.MakeKeepClient(&arv)
  86         if err != nil {
  87                 log.Fatalf("Error setting up keep client %s", err.Error())
  88         }
  89
  90         if pidfile != "" {
  91                 f, err := os.Create(pidfile)
  92                 if err != nil {
  93                         log.Fatalf("Error writing pid file (%s): %s", pidfile, err.Error())
  94                 }
  95                 fmt.Fprint(f, os.Getpid())
  96                 f.Close()
  97                 defer os.Remove(pidfile)
  98         }
  99
 100         kc.Want_replicas = default_replicas
 101
 102         kc.Client.Timeout = time.Duration(timeout) * time.Second
 103
 104         listener, err = net.Listen("tcp", listen)
 105         if err != nil {
 106                 log.Fatalf("Could not listen on %v", listen)
 107         }
 108
 109         go RefreshServicesList(kc)
 110
 111         // Shut down the server gracefully (by closing the listener)
 112         // if SIGTERM is received.
 113         term := make(chan os.Signal, 1)
 114         go func(sig <-chan os.Signal) {
 115                 s := <-sig
 116                 log.Println("caught signal:", s)
 117                 listener.Close()
 118         }(term)
 119         signal.Notify(term, syscall.SIGTERM)
 120         signal.Notify(term, syscall.SIGINT)
 121
 122         log.Printf("Arvados Keep proxy started listening on %v", listener.Addr())
 123
 124         // Start listening for requests.
 125         http.Serve(listener, MakeRESTRouter(!no_get, !no_put, kc))
 126
 127         log.Println("shutting down")
 128 }
 129
 130 type ApiTokenCache struct {
 131         tokens     map[string]int64
 132         lock       sync.Mutex
 133         expireTime int64
 134 }
 135
 136 // Refresh the keep service list every five minutes.
 137 func RefreshServicesList(kc *keepclient.KeepClient) {
 138         previousRoots := ""
 139         for {
 140                 if err := kc.DiscoverKeepServers(); err != nil {
 141                         log.Println("Error retrieving services list:", err)
 142                         time.Sleep(3*time.Second)
 143                         previousRoots = ""
 144                 } else if len(kc.LocalRoots()) == 0 {
 145                         log.Println("Received empty services list")
 146                         time.Sleep(3*time.Second)
 147                         previousRoots = ""
 148                 } else {
 149                         newRoots := fmt.Sprint("Locals ", kc.LocalRoots(), ", gateways ", kc.GatewayRoots())
 150                         if newRoots != previousRoots {
 151                                 log.Println("Updated services list:", newRoots)
 152                                 previousRoots = newRoots
 153                         }
 154                         time.Sleep(300*time.Second)
 155                 }
 156         }
 157 }
 158
 159 // Cache the token and set an expire time.  If we already have an expire time
 160 // on the token, it is not updated.
 161 func (this *ApiTokenCache) RememberToken(token string) {
 162         this.lock.Lock()
 163         defer this.lock.Unlock()
 164
 165         now := time.Now().Unix()
 166         if this.tokens[token] == 0 {
 167                 this.tokens[token] = now + this.expireTime
 168         }
 169 }
 170
 171 // Check if the cached token is known and still believed to be valid.
 172 func (this *ApiTokenCache) RecallToken(token string) bool {
 173         this.lock.Lock()
 174         defer this.lock.Unlock()
 175
 176         now := time.Now().Unix()
 177         if this.tokens[token] == 0 {
 178                 // Unknown token
 179                 return false
 180         } else if now < this.tokens[token] {
 181                 // Token is known and still valid
 182                 return true
 183         } else {
 184                 // Token is expired
 185                 this.tokens[token] = 0
 186                 return false
 187         }
 188 }
 189
 190 func GetRemoteAddress(req *http.Request) string {
 191         if realip := req.Header.Get("X-Real-IP"); realip != "" {
 192                 if forwarded := req.Header.Get("X-Forwarded-For"); forwarded != realip {
 193                         return fmt.Sprintf("%s (X-Forwarded-For %s)", realip, forwarded)
 194                 } else {
 195                         return realip
 196                 }
 197         }
 198         return req.RemoteAddr
 199 }
 200
 201 func CheckAuthorizationHeader(kc keepclient.KeepClient, cache *ApiTokenCache, req *http.Request) (pass bool, tok string) {
 202         var auth string
 203         if auth = req.Header.Get("Authorization"); auth == "" {
 204                 return false, ""
 205         }
 206
 207         _, err := fmt.Sscanf(auth, "OAuth2 %s", &tok)
 208         if err != nil {
 209                 // Scanning error
 210                 return false, ""
 211         }
 212
 213         if cache.RecallToken(tok) {
 214                 // Valid in the cache, short circut
 215                 return true, tok
 216         }
 217
 218         arv := *kc.Arvados
 219         arv.ApiToken = tok
 220         if err := arv.Call("HEAD", "users", "", "current", nil, nil); err != nil {
 221                 log.Printf("%s: CheckAuthorizationHeader error: %v", GetRemoteAddress(req), err)
 222                 return false, ""
 223         }
 224
 225         // Success!  Update cache
 226         cache.RememberToken(tok)
 227
 228         return true, tok
 229 }
 230
 231 type GetBlockHandler struct {
 232         *keepclient.KeepClient
 233         *ApiTokenCache
 234 }
 235
 236 type PutBlockHandler struct {
 237         *keepclient.KeepClient
 238         *ApiTokenCache
 239 }
 240
 241 type InvalidPathHandler struct{}
 242
 243 type OptionsHandler struct{}
 244
 245 // MakeRESTRouter
 246 //     Returns a mux.Router that passes GET and PUT requests to the
 247 //     appropriate handlers.
 248 //
 249 func MakeRESTRouter(
 250         enable_get bool,
 251         enable_put bool,
 252         kc *keepclient.KeepClient) *mux.Router {
 253
 254         t := &ApiTokenCache{tokens: make(map[string]int64), expireTime: 300}
 255
 256         rest := mux.NewRouter()
 257
 258         if enable_get {
 259                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`,
 260                         GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 261                 rest.Handle(`/{locator:[0-9a-f]{32}}`, GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 262         }
 263
 264         if enable_put {
 265                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`, PutBlockHandler{kc, t}).Methods("PUT")
 266                 rest.Handle(`/{locator:[0-9a-f]{32}}`, PutBlockHandler{kc, t}).Methods("PUT")
 267                 rest.Handle(`/`, PutBlockHandler{kc, t}).Methods("POST")
 268                 rest.Handle(`/{any}`, OptionsHandler{}).Methods("OPTIONS")
 269                 rest.Handle(`/`, OptionsHandler{}).Methods("OPTIONS")
 270         }
 271
 272         rest.NotFoundHandler = InvalidPathHandler{}
 273
 274         return rest
 275 }
 276
 277 func SetCorsHeaders(resp http.ResponseWriter) {
 278         resp.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, POST, PUT, OPTIONS")
 279         resp.Header().Set("Access-Control-Allow-Origin", "*")
 280         resp.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Length, Content-Type, X-Keep-Desired-Replicas")
 281         resp.Header().Set("Access-Control-Max-Age", "86486400")
 282 }
 283
 284 func (this InvalidPathHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 285         log.Printf("%s: %s %s unroutable", GetRemoteAddress(req), req.Method, req.URL.Path)
 286         http.Error(resp, "Bad request", http.StatusBadRequest)
 287 }
 288
 289 func (this OptionsHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 290         log.Printf("%s: %s %s", GetRemoteAddress(req), req.Method, req.URL.Path)
 291         SetCorsHeaders(resp)
 292 }
 293
 294 var BadAuthorizationHeader = errors.New("Missing or invalid Authorization header")
 295 var ContentLengthMismatch = errors.New("Actual length != expected content length")
 296 var MethodNotSupported = errors.New("Method not supported")
 297
 298 func (this GetBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 299         SetCorsHeaders(resp)
 300
 301         locator := mux.Vars(req)["locator"]
 302         var err error
 303         var status int
 304         var expectLength, responseLength int64
 305         var proxiedURI = "-"
 306
 307         defer func() {
 308                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, responseLength, proxiedURI, err)
 309                 if status != http.StatusOK {
 310                         http.Error(resp, err.Error(), status)
 311                 }
 312         }()
 313
 314         kc := *this.KeepClient
 315
 316         var pass bool
 317         var tok string
 318         if pass, tok = CheckAuthorizationHeader(kc, this.ApiTokenCache, req); !pass {
 319                 status, err = http.StatusForbidden, BadAuthorizationHeader
 320                 return
 321         }
 322
 323         // Copy ArvadosClient struct and use the client's API token
 324         arvclient := *kc.Arvados
 325         arvclient.ApiToken = tok
 326         kc.Arvados = &arvclient
 327
 328         var reader io.ReadCloser
 329
 330         switch req.Method {
 331         case "HEAD":
 332                 expectLength, proxiedURI, err = kc.Ask(locator)
 333         case "GET":
 334                 reader, expectLength, proxiedURI, err = kc.Get(locator)
 335                 if reader != nil {
 336                         defer reader.Close()
 337                 }
 338         default:
 339                 status, err = http.StatusNotImplemented, MethodNotSupported
 340                 return
 341         }
 342
 343         if expectLength == -1 {
 344                 log.Println("Warning:", GetRemoteAddress(req), req.Method, proxiedURI, "Content-Length not provided")
 345         }
 346
 347         switch err {
 348         case nil:
 349                 status = http.StatusOK
 350                 resp.Header().Set("Content-Length", fmt.Sprint(expectLength))
 351                 switch req.Method {
 352                 case "HEAD":
 353                         responseLength = 0
 354                 case "GET":
 355                         responseLength, err = io.Copy(resp, reader)
 356                         if err == nil && expectLength > -1 && responseLength != expectLength {
 357                                 err = ContentLengthMismatch
 358                         }
 359                 }
 360         case keepclient.BlockNotFound:
 361                 status = http.StatusNotFound
 362         default:
 363                 status = http.StatusBadGateway
 364         }
 365 }
 366
 367 var LengthRequiredError = errors.New(http.StatusText(http.StatusLengthRequired))
 368 var LengthMismatchError = errors.New("Locator size hint does not match Content-Length header")
 369
 370 func (this PutBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 371         SetCorsHeaders(resp)
 372
 373         kc := *this.KeepClient
 374         var err error
 375         var expectLength int64 = -1
 376         var status = http.StatusInternalServerError
 377         var wroteReplicas int
 378         var locatorOut string = "-"
 379
 380         defer func() {
 381                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, kc.Want_replicas, wroteReplicas, locatorOut, err)
 382                 if status != http.StatusOK {
 383                         http.Error(resp, err.Error(), status)
 384                 }
 385         }()
 386
 387         locatorIn := mux.Vars(req)["locator"]
 388
 389         if req.Header.Get("Content-Length") != "" {
 390                 _, err := fmt.Sscanf(req.Header.Get("Content-Length"), "%d", &expectLength)
 391                 if err != nil {
 392                         resp.Header().Set("Content-Length", fmt.Sprintf("%d", expectLength))
 393                 }
 394
 395         }
 396
 397         if expectLength < 0 {
 398                 err = LengthRequiredError
 399                 status = http.StatusLengthRequired
 400                 return
 401         }
 402
 403         if locatorIn != "" {
 404                 var loc *keepclient.Locator
 405                 if loc, err = keepclient.MakeLocator(locatorIn); err != nil {
 406                         status = http.StatusBadRequest
 407                         return
 408                 } else if loc.Size > 0 && int64(loc.Size) != expectLength {
 409                         err = LengthMismatchError
 410                         status = http.StatusBadRequest
 411                         return
 412                 }
 413         }
 414
 415         var pass bool
 416         var tok string
 417         if pass, tok = CheckAuthorizationHeader(kc, this.ApiTokenCache, req); !pass {
 418                 err = BadAuthorizationHeader
 419                 status = http.StatusForbidden
 420                 return
 421         }
 422
 423         // Copy ArvadosClient struct and use the client's API token
 424         arvclient := *kc.Arvados
 425         arvclient.ApiToken = tok
 426         kc.Arvados = &arvclient
 427
 428         // Check if the client specified the number of replicas
 429         if req.Header.Get("X-Keep-Desired-Replicas") != "" {
 430                 var r int
 431                 _, err := fmt.Sscanf(req.Header.Get(keepclient.X_Keep_Desired_Replicas), "%d", &r)
 432                 if err != nil {
 433                         kc.Want_replicas = r
 434                 }
 435         }
 436
 437         // Now try to put the block through
 438         if locatorIn == "" {
 439                 if bytes, err := ioutil.ReadAll(req.Body); err != nil {
 440                         err = errors.New(fmt.Sprintf("Error reading request body: %s", err))
 441                         status = http.StatusInternalServerError
 442                         return
 443                 } else {
 444                         locatorOut, wroteReplicas, err = kc.PutB(bytes)
 445                 }
 446         } else {
 447                 locatorOut, wroteReplicas, err = kc.PutHR(locatorIn, req.Body, expectLength)
 448         }
 449
 450         // Tell the client how many successful PUTs we accomplished
 451         resp.Header().Set(keepclient.X_Keep_Replicas_Stored, fmt.Sprintf("%d", wroteReplicas))
 452
 453         switch err {
 454         case nil:
 455                 status = http.StatusOK
 456                 _, err = io.WriteString(resp, locatorOut)
 457
 458         case keepclient.OversizeBlockError:
 459                 // Too much data
 460                 status = http.StatusRequestEntityTooLarge
 461
 462         case keepclient.InsufficientReplicasError:
 463                 if wroteReplicas > 0 {
 464                         // At least one write is considered success.  The
 465                         // client can decide if getting less than the number of
 466                         // replications it asked for is a fatal error.
 467                         status = http.StatusOK
 468                         _, err = io.WriteString(resp, locatorOut)
 469                 } else {
 470                         status = http.StatusServiceUnavailable
 471                 }
 472
 473         default:
 474                 status = http.StatusBadGateway
 475         }
 476 }