services/keepproxy/keepproxy.go

   1 package main
   2
   3 import (
   4         "errors"
   5         "flag"
   6         "fmt"
   7         "io"
   8         "io/ioutil"
   9         "log"
  10         "net"
  11         "net/http"
  12         "os"
  13         "os/signal"
  14         "regexp"
  15         "sync"
  16         "syscall"
  17         "time"
  18
  19         "git.curoverse.com/arvados.git/sdk/go/arvados"
  20         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
  21         "git.curoverse.com/arvados.git/sdk/go/config"
  22         "git.curoverse.com/arvados.git/sdk/go/keepclient"
  23         "github.com/coreos/go-systemd/daemon"
  24         "github.com/ghodss/yaml"
  25         "github.com/gorilla/mux"
  26 )
  27
  28 type Config struct {
  29         Client          arvados.Client
  30         Listen          string
  31         DisableGet      bool
  32         DisablePut      bool
  33         DefaultReplicas int
  34         Timeout         arvados.Duration
  35         PIDFile         string
  36         Debug           bool
  37 }
  38
  39 func DefaultConfig() *Config {
  40         return &Config{
  41                 Listen:  ":25107",
  42                 Timeout: arvados.Duration(15 * time.Second),
  43         }
  44 }
  45
  46 var listener net.Listener
  47
  48 func main() {
  49         cfg := DefaultConfig()
  50
  51         flagset := flag.NewFlagSet("keepproxy", flag.ExitOnError)
  52         flagset.Usage = usage
  53
  54         const deprecated = " (DEPRECATED -- use config file instead)"
  55         flagset.StringVar(&cfg.Listen, "listen", cfg.Listen, "Local port to listen on."+deprecated)
  56         flagset.BoolVar(&cfg.DisableGet, "no-get", cfg.DisableGet, "Disable GET operations."+deprecated)
  57         flagset.BoolVar(&cfg.DisablePut, "no-put", cfg.DisablePut, "Disable PUT operations."+deprecated)
  58         flagset.IntVar(&cfg.DefaultReplicas, "default-replicas", cfg.DefaultReplicas, "Default number of replicas to write if not specified by the client. If 0, use site default."+deprecated)
  59         flagset.StringVar(&cfg.PIDFile, "pid", cfg.PIDFile, "Path to write pid file."+deprecated)
  60         timeoutSeconds := flagset.Int("timeout", int(time.Duration(cfg.Timeout)/time.Second), "Timeout (in seconds) on requests to internal Keep services."+deprecated)
  61
  62         var cfgPath string
  63         const defaultCfgPath = "/etc/arvados/keepproxy/keepproxy.yml"
  64         flagset.StringVar(&cfgPath, "config", defaultCfgPath, "Configuration file `path`")
  65         dumpConfig := flagset.Bool("dump-config", false, "write current configuration to stdout and exit")
  66         flagset.Parse(os.Args[1:])
  67
  68         err := config.LoadFile(cfg, cfgPath)
  69         if err != nil {
  70                 h := os.Getenv("ARVADOS_API_HOST")
  71                 t := os.Getenv("ARVADOS_API_TOKEN")
  72                 if h == "" || t == "" || !os.IsNotExist(err) || cfgPath != defaultCfgPath {
  73                         log.Fatal(err)
  74                 }
  75                 log.Print("DEPRECATED: No config file found, but ARVADOS_API_HOST and ARVADOS_API_TOKEN environment variables are set. Please use a config file instead.")
  76                 cfg.Client.APIHost = h
  77                 cfg.Client.AuthToken = t
  78                 if regexp.MustCompile("^(?i:1|yes|true)$").MatchString(os.Getenv("ARVADOS_API_HOST_INSECURE")) {
  79                         cfg.Client.Insecure = true
  80                 }
  81                 if y, err := yaml.Marshal(cfg); err == nil && !*dumpConfig {
  82                         log.Print("Current configuration:\n", string(y))
  83                 }
  84                 cfg.Timeout = arvados.Duration(time.Duration(*timeoutSeconds) * time.Second)
  85         }
  86
  87         if *dumpConfig {
  88                 log.Fatal(config.DumpAndExit(cfg))
  89         }
  90
  91         arv, err := arvadosclient.New(&cfg.Client)
  92         if err != nil {
  93                 log.Fatalf("Error setting up arvados client %s", err.Error())
  94         }
  95
  96         if cfg.Debug {
  97                 keepclient.DebugPrintf = log.Printf
  98         }
  99         kc, err := keepclient.MakeKeepClient(arv)
 100         if err != nil {
 101                 log.Fatalf("Error setting up keep client %s", err.Error())
 102         }
 103
 104         if cfg.PIDFile != "" {
 105                 f, err := os.Create(cfg.PIDFile)
 106                 if err != nil {
 107                         log.Fatal(err)
 108                 }
 109                 defer f.Close()
 110                 err = syscall.Flock(int(f.Fd()), syscall.LOCK_EX|syscall.LOCK_NB)
 111                 if err != nil {
 112                         log.Fatalf("flock(%s): %s", cfg.PIDFile, err)
 113                 }
 114                 defer os.Remove(cfg.PIDFile)
 115                 err = f.Truncate(0)
 116                 if err != nil {
 117                         log.Fatalf("truncate(%s): %s", cfg.PIDFile, err)
 118                 }
 119                 _, err = fmt.Fprint(f, os.Getpid())
 120                 if err != nil {
 121                         log.Fatalf("write(%s): %s", cfg.PIDFile, err)
 122                 }
 123                 err = f.Sync()
 124                 if err != nil {
 125                         log.Fatal("sync(%s): %s", cfg.PIDFile, err)
 126                 }
 127         }
 128
 129         if cfg.DefaultReplicas > 0 {
 130                 kc.Want_replicas = cfg.DefaultReplicas
 131         }
 132         kc.Client.Timeout = time.Duration(cfg.Timeout)
 133         go kc.RefreshServices(5*time.Minute, 3*time.Second)
 134
 135         listener, err = net.Listen("tcp", cfg.Listen)
 136         if err != nil {
 137                 log.Fatalf("listen(%s): %s", cfg.Listen, err)
 138         }
 139         if _, err := daemon.SdNotify(false, "READY=1"); err != nil {
 140                 log.Printf("Error notifying init daemon: %v", err)
 141         }
 142         log.Println("Listening at", listener.Addr())
 143
 144         // Shut down the server gracefully (by closing the listener)
 145         // if SIGTERM is received.
 146         term := make(chan os.Signal, 1)
 147         go func(sig <-chan os.Signal) {
 148                 s := <-sig
 149                 log.Println("caught signal:", s)
 150                 listener.Close()
 151         }(term)
 152         signal.Notify(term, syscall.SIGTERM)
 153         signal.Notify(term, syscall.SIGINT)
 154
 155         // Start serving requests.
 156         http.Serve(listener, MakeRESTRouter(!cfg.DisableGet, !cfg.DisablePut, kc))
 157
 158         log.Println("shutting down")
 159 }
 160
 161 type ApiTokenCache struct {
 162         tokens     map[string]int64
 163         lock       sync.Mutex
 164         expireTime int64
 165 }
 166
 167 // Cache the token and set an expire time.  If we already have an expire time
 168 // on the token, it is not updated.
 169 func (this *ApiTokenCache) RememberToken(token string) {
 170         this.lock.Lock()
 171         defer this.lock.Unlock()
 172
 173         now := time.Now().Unix()
 174         if this.tokens[token] == 0 {
 175                 this.tokens[token] = now + this.expireTime
 176         }
 177 }
 178
 179 // Check if the cached token is known and still believed to be valid.
 180 func (this *ApiTokenCache) RecallToken(token string) bool {
 181         this.lock.Lock()
 182         defer this.lock.Unlock()
 183
 184         now := time.Now().Unix()
 185         if this.tokens[token] == 0 {
 186                 // Unknown token
 187                 return false
 188         } else if now < this.tokens[token] {
 189                 // Token is known and still valid
 190                 return true
 191         } else {
 192                 // Token is expired
 193                 this.tokens[token] = 0
 194                 return false
 195         }
 196 }
 197
 198 func GetRemoteAddress(req *http.Request) string {
 199         if xff := req.Header.Get("X-Forwarded-For"); xff != "" {
 200                 return xff + "," + req.RemoteAddr
 201         }
 202         return req.RemoteAddr
 203 }
 204
 205 func CheckAuthorizationHeader(kc *keepclient.KeepClient, cache *ApiTokenCache, req *http.Request) (pass bool, tok string) {
 206         var auth string
 207         if auth = req.Header.Get("Authorization"); auth == "" {
 208                 return false, ""
 209         }
 210
 211         _, err := fmt.Sscanf(auth, "OAuth2 %s", &tok)
 212         if err != nil {
 213                 // Scanning error
 214                 return false, ""
 215         }
 216
 217         if cache.RecallToken(tok) {
 218                 // Valid in the cache, short circuit
 219                 return true, tok
 220         }
 221
 222         arv := *kc.Arvados
 223         arv.ApiToken = tok
 224         if err := arv.Call("HEAD", "users", "", "current", nil, nil); err != nil {
 225                 log.Printf("%s: CheckAuthorizationHeader error: %v", GetRemoteAddress(req), err)
 226                 return false, ""
 227         }
 228
 229         // Success!  Update cache
 230         cache.RememberToken(tok)
 231
 232         return true, tok
 233 }
 234
 235 type GetBlockHandler struct {
 236         *keepclient.KeepClient
 237         *ApiTokenCache
 238 }
 239
 240 type PutBlockHandler struct {
 241         *keepclient.KeepClient
 242         *ApiTokenCache
 243 }
 244
 245 type IndexHandler struct {
 246         *keepclient.KeepClient
 247         *ApiTokenCache
 248 }
 249
 250 type InvalidPathHandler struct{}
 251
 252 type OptionsHandler struct{}
 253
 254 // MakeRESTRouter
 255 //     Returns a mux.Router that passes GET and PUT requests to the
 256 //     appropriate handlers.
 257 //
 258 func MakeRESTRouter(
 259         enable_get bool,
 260         enable_put bool,
 261         kc *keepclient.KeepClient) *mux.Router {
 262
 263         t := &ApiTokenCache{tokens: make(map[string]int64), expireTime: 300}
 264
 265         rest := mux.NewRouter()
 266
 267         if enable_get {
 268                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`,
 269                         GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 270                 rest.Handle(`/{locator:[0-9a-f]{32}}`, GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 271
 272                 // List all blocks
 273                 rest.Handle(`/index`, IndexHandler{kc, t}).Methods("GET")
 274
 275                 // List blocks whose hash has the given prefix
 276                 rest.Handle(`/index/{prefix:[0-9a-f]{0,32}}`, IndexHandler{kc, t}).Methods("GET")
 277         }
 278
 279         if enable_put {
 280                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`, PutBlockHandler{kc, t}).Methods("PUT")
 281                 rest.Handle(`/{locator:[0-9a-f]{32}}`, PutBlockHandler{kc, t}).Methods("PUT")
 282                 rest.Handle(`/`, PutBlockHandler{kc, t}).Methods("POST")
 283                 rest.Handle(`/{any}`, OptionsHandler{}).Methods("OPTIONS")
 284                 rest.Handle(`/`, OptionsHandler{}).Methods("OPTIONS")
 285         }
 286
 287         rest.NotFoundHandler = InvalidPathHandler{}
 288
 289         return rest
 290 }
 291
 292 func SetCorsHeaders(resp http.ResponseWriter) {
 293         resp.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, POST, PUT, OPTIONS")
 294         resp.Header().Set("Access-Control-Allow-Origin", "*")
 295         resp.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Length, Content-Type, X-Keep-Desired-Replicas")
 296         resp.Header().Set("Access-Control-Max-Age", "86486400")
 297 }
 298
 299 func (this InvalidPathHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 300         log.Printf("%s: %s %s unroutable", GetRemoteAddress(req), req.Method, req.URL.Path)
 301         http.Error(resp, "Bad request", http.StatusBadRequest)
 302 }
 303
 304 func (this OptionsHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 305         log.Printf("%s: %s %s", GetRemoteAddress(req), req.Method, req.URL.Path)
 306         SetCorsHeaders(resp)
 307 }
 308
 309 var BadAuthorizationHeader = errors.New("Missing or invalid Authorization header")
 310 var ContentLengthMismatch = errors.New("Actual length != expected content length")
 311 var MethodNotSupported = errors.New("Method not supported")
 312
 313 var removeHint, _ = regexp.Compile("\\+K@[a-z0-9]{5}(\\+|$)")
 314
 315 func (this GetBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 316         SetCorsHeaders(resp)
 317
 318         locator := mux.Vars(req)["locator"]
 319         var err error
 320         var status int
 321         var expectLength, responseLength int64
 322         var proxiedURI = "-"
 323
 324         defer func() {
 325                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, responseLength, proxiedURI, err)
 326                 if status != http.StatusOK {
 327                         http.Error(resp, err.Error(), status)
 328                 }
 329         }()
 330
 331         kc := *this.KeepClient
 332
 333         var pass bool
 334         var tok string
 335         if pass, tok = CheckAuthorizationHeader(&kc, this.ApiTokenCache, req); !pass {
 336                 status, err = http.StatusForbidden, BadAuthorizationHeader
 337                 return
 338         }
 339
 340         // Copy ArvadosClient struct and use the client's API token
 341         arvclient := *kc.Arvados
 342         arvclient.ApiToken = tok
 343         kc.Arvados = &arvclient
 344
 345         var reader io.ReadCloser
 346
 347         locator = removeHint.ReplaceAllString(locator, "$1")
 348
 349         switch req.Method {
 350         case "HEAD":
 351                 expectLength, proxiedURI, err = kc.Ask(locator)
 352         case "GET":
 353                 reader, expectLength, proxiedURI, err = kc.Get(locator)
 354                 if reader != nil {
 355                         defer reader.Close()
 356                 }
 357         default:
 358                 status, err = http.StatusNotImplemented, MethodNotSupported
 359                 return
 360         }
 361
 362         if expectLength == -1 {
 363                 log.Println("Warning:", GetRemoteAddress(req), req.Method, proxiedURI, "Content-Length not provided")
 364         }
 365
 366         switch respErr := err.(type) {
 367         case nil:
 368                 status = http.StatusOK
 369                 resp.Header().Set("Content-Length", fmt.Sprint(expectLength))
 370                 switch req.Method {
 371                 case "HEAD":
 372                         responseLength = 0
 373                 case "GET":
 374                         responseLength, err = io.Copy(resp, reader)
 375                         if err == nil && expectLength > -1 && responseLength != expectLength {
 376                                 err = ContentLengthMismatch
 377                         }
 378                 }
 379         case keepclient.Error:
 380                 if respErr == keepclient.BlockNotFound {
 381                         status = http.StatusNotFound
 382                 } else if respErr.Temporary() {
 383                         status = http.StatusBadGateway
 384                 } else {
 385                         status = 422
 386                 }
 387         default:
 388                 status = http.StatusInternalServerError
 389         }
 390 }
 391
 392 var LengthRequiredError = errors.New(http.StatusText(http.StatusLengthRequired))
 393 var LengthMismatchError = errors.New("Locator size hint does not match Content-Length header")
 394
 395 func (this PutBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 396         SetCorsHeaders(resp)
 397
 398         kc := *this.KeepClient
 399         var err error
 400         var expectLength int64
 401         var status = http.StatusInternalServerError
 402         var wroteReplicas int
 403         var locatorOut string = "-"
 404
 405         defer func() {
 406                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, kc.Want_replicas, wroteReplicas, locatorOut, err)
 407                 if status != http.StatusOK {
 408                         http.Error(resp, err.Error(), status)
 409                 }
 410         }()
 411
 412         locatorIn := mux.Vars(req)["locator"]
 413
 414         _, err = fmt.Sscanf(req.Header.Get("Content-Length"), "%d", &expectLength)
 415         if err != nil || expectLength < 0 {
 416                 err = LengthRequiredError
 417                 status = http.StatusLengthRequired
 418                 return
 419         }
 420
 421         if locatorIn != "" {
 422                 var loc *keepclient.Locator
 423                 if loc, err = keepclient.MakeLocator(locatorIn); err != nil {
 424                         status = http.StatusBadRequest
 425                         return
 426                 } else if loc.Size > 0 && int64(loc.Size) != expectLength {
 427                         err = LengthMismatchError
 428                         status = http.StatusBadRequest
 429                         return
 430                 }
 431         }
 432
 433         var pass bool
 434         var tok string
 435         if pass, tok = CheckAuthorizationHeader(&kc, this.ApiTokenCache, req); !pass {
 436                 err = BadAuthorizationHeader
 437                 status = http.StatusForbidden
 438                 return
 439         }
 440
 441         // Copy ArvadosClient struct and use the client's API token
 442         arvclient := *kc.Arvados
 443         arvclient.ApiToken = tok
 444         kc.Arvados = &arvclient
 445
 446         // Check if the client specified the number of replicas
 447         if req.Header.Get("X-Keep-Desired-Replicas") != "" {
 448                 var r int
 449                 _, err := fmt.Sscanf(req.Header.Get(keepclient.X_Keep_Desired_Replicas), "%d", &r)
 450                 if err == nil {
 451                         kc.Want_replicas = r
 452                 }
 453         }
 454
 455         // Now try to put the block through
 456         if locatorIn == "" {
 457                 if bytes, err := ioutil.ReadAll(req.Body); err != nil {
 458                         err = errors.New(fmt.Sprintf("Error reading request body: %s", err))
 459                         status = http.StatusInternalServerError
 460                         return
 461                 } else {
 462                         locatorOut, wroteReplicas, err = kc.PutB(bytes)
 463                 }
 464         } else {
 465                 locatorOut, wroteReplicas, err = kc.PutHR(locatorIn, req.Body, expectLength)
 466         }
 467
 468         // Tell the client how many successful PUTs we accomplished
 469         resp.Header().Set(keepclient.X_Keep_Replicas_Stored, fmt.Sprintf("%d", wroteReplicas))
 470
 471         switch err {
 472         case nil:
 473                 status = http.StatusOK
 474                 _, err = io.WriteString(resp, locatorOut)
 475
 476         case keepclient.OversizeBlockError:
 477                 // Too much data
 478                 status = http.StatusRequestEntityTooLarge
 479
 480         case keepclient.InsufficientReplicasError:
 481                 if wroteReplicas > 0 {
 482                         // At least one write is considered success.  The
 483                         // client can decide if getting less than the number of
 484                         // replications it asked for is a fatal error.
 485                         status = http.StatusOK
 486                         _, err = io.WriteString(resp, locatorOut)
 487                 } else {
 488                         status = http.StatusServiceUnavailable
 489                 }
 490
 491         default:
 492                 status = http.StatusBadGateway
 493         }
 494 }
 495
 496 // ServeHTTP implementation for IndexHandler
 497 // Supports only GET requests for /index/{prefix:[0-9a-f]{0,32}}
 498 // For each keep server found in LocalRoots:
 499 //   Invokes GetIndex using keepclient
 500 //   Expects "complete" response (terminating with blank new line)
 501 //   Aborts on any errors
 502 // Concatenates responses from all those keep servers and returns
 503 func (handler IndexHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 504         SetCorsHeaders(resp)
 505
 506         prefix := mux.Vars(req)["prefix"]
 507         var err error
 508         var status int
 509
 510         defer func() {
 511                 if status != http.StatusOK {
 512                         http.Error(resp, err.Error(), status)
 513                 }
 514         }()
 515
 516         kc := *handler.KeepClient
 517
 518         ok, token := CheckAuthorizationHeader(&kc, handler.ApiTokenCache, req)
 519         if !ok {
 520                 status, err = http.StatusForbidden, BadAuthorizationHeader
 521                 return
 522         }
 523
 524         // Copy ArvadosClient struct and use the client's API token
 525         arvclient := *kc.Arvados
 526         arvclient.ApiToken = token
 527         kc.Arvados = &arvclient
 528
 529         // Only GET method is supported
 530         if req.Method != "GET" {
 531                 status, err = http.StatusNotImplemented, MethodNotSupported
 532                 return
 533         }
 534
 535         // Get index from all LocalRoots and write to resp
 536         var reader io.Reader
 537         for uuid := range kc.LocalRoots() {
 538                 reader, err = kc.GetIndex(uuid, prefix)
 539                 if err != nil {
 540                         status = http.StatusBadGateway
 541                         return
 542                 }
 543
 544                 _, err = io.Copy(resp, reader)
 545                 if err != nil {
 546                         status = http.StatusBadGateway
 547                         return
 548                 }
 549         }
 550
 551         // Got index from all the keep servers and wrote to resp
 552         status = http.StatusOK
 553         resp.Write([]byte("\n"))
 554 }