services/keepstore/pull_worker.go

   1 package main
   2
   3 import (
   4         "crypto/rand"
   5         "errors"
   6         "fmt"
   7         "git.curoverse.com/arvados.git/sdk/go/keepclient"
   8         "io"
   9         "io/ioutil"
  10         "log"
  11         "os"
  12         "time"
  13 )
  14
  15 var keepClient keepclient.KeepClient
  16
  17 /*
  18         Keepstore initiates pull worker channel goroutine.
  19         The channel will process pull list.
  20                 For each (next) pull request:
  21                         For each locator listed, execute Pull on the server(s) listed
  22                         Skip the rest of the servers if no errors
  23                 Repeat
  24 */
  25 func RunPullWorker(pullq *WorkQueue, kc keepclient.KeepClient) {
  26         nextItem := pullq.NextItem
  27         keepClient = kc
  28         for item := range nextItem {
  29                 Pull(item.(PullRequest))
  30         }
  31 }
  32
  33 /*
  34         For each Pull request:
  35                 Generate a random API token.
  36                 Generate a permission signature using this token, timestamp ~60 seconds in the future, and desired block hash.
  37                 Using this token & signature, retrieve the given block.
  38                 Write to storage
  39 */
  40 func Pull(pullRequest PullRequest) (err error) {
  41         defer func() {
  42                 if err == nil {
  43                         log.Printf("Pull %s success", pullRequest)
  44                 } else {
  45                         log.Printf("Pull %s error: %s", pullRequest, err)
  46                 }
  47         }()
  48
  49         service_roots := make(map[string]string)
  50         for _, addr := range pullRequest.Servers {
  51                 service_roots[addr] = addr
  52         }
  53         keepClient.SetServiceRoots(service_roots)
  54
  55         // Generate signature with a random token
  56         PermissionSecret = []byte(os.Getenv("ARVADOS_API_TOKEN"))
  57         expires_at := time.Now().Add(60 * time.Second)
  58         signedLocator := SignLocator(pullRequest.Locator, GenerateRandomApiToken(), expires_at)
  59
  60         reader, contentLen, _, err := GetContent(pullRequest.Locator, signedLocator)
  61         if err != nil {
  62                 return
  63         }
  64         if reader == nil {
  65                 return errors.New(fmt.Sprintf("No reader found for : %s", signedLocator))
  66         }
  67         defer reader.Close()
  68
  69         read_content, err := ioutil.ReadAll(reader)
  70         if err != nil {
  71                 return err
  72         }
  73
  74         if (read_content == nil) || (int64(len(read_content)) != contentLen) {
  75                 return errors.New(fmt.Sprintf("Content not found for: %s", signedLocator))
  76         }
  77
  78         err = PutContent(read_content, pullRequest.Locator)
  79         return
  80 }
  81
  82 // Fetch the content for the given locator using keepclient.
  83 var GetContent = func(locator string, signedLocator string) (reader io.ReadCloser, contentLength int64, url string, err error) {
  84         reader, blocklen, url, err := keepClient.Get(signedLocator)
  85         return reader, blocklen, url, err
  86 }
  87
  88 const ALPHA_NUMERIC = "0123456789abcdefghijklmnopqrstuvwxyz"
  89
  90 func GenerateRandomApiToken() string {
  91         var bytes = make([]byte, 36)
  92         rand.Read(bytes)
  93         for i, b := range bytes {
  94                 bytes[i] = ALPHA_NUMERIC[b%byte(len(ALPHA_NUMERIC))]
  95         }
  96         return (string(bytes))
  97 }
  98
  99 // Put block
 100 var PutContent = func(content []byte, locator string) (err error) {
 101         err = PutBlock(content, locator)
 102         return
 103 }