X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/bbb0760e909b3ceca850b3aa319011fba2c98ed0..77c8223f5ddd64cff2b08d0857749644c474946f:/tools/keep-block-check/keep-block-check.go

diff --git a/tools/keep-block-check/keep-block-check.go b/tools/keep-block-check/keep-block-check.go
index caed9d0469..60d72773c1 100644
--- a/tools/keep-block-check/keep-block-check.go
+++ b/tools/keep-block-check/keep-block-check.go
@@ -1,3 +1,7 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package main
 
 import (
@@ -5,25 +9,27 @@ import (
 	"errors"
 	"flag"
 	"fmt"
-	"git.curoverse.com/arvados.git/sdk/go/arvadosclient"
-	"git.curoverse.com/arvados.git/sdk/go/keepclient"
 	"io/ioutil"
 	"log"
 	"net/http"
 	"os"
-	"regexp"
 	"strings"
 	"time"
+
+	"git.arvados.org/arvados.git/sdk/go/arvadosclient"
+	"git.arvados.org/arvados.git/sdk/go/keepclient"
 )
 
+var version = "dev"
+
 func main() {
-	err := doMain()
+	err := doMain(os.Args[1:])
 	if err != nil {
 		log.Fatalf("%v", err)
 	}
 }
 
-func doMain() error {
+func doMain(args []string) error {
 	flags := flag.NewFlagSet("keep-block-check", flag.ExitOnError)
 
 	configFile := flags.String(
@@ -48,8 +54,29 @@ func doMain() error {
 		"",
 		"Block hash prefix. When a prefix is specified, only hashes listed in the file with this prefix will be checked.")
 
+	blobSignatureTTLFlag := flags.Duration(
+		"blob-signature-ttl",
+		0,
+		"Lifetime of blob permission signatures on the keepservers. If not provided, this will be retrieved from the API server's discovery document.")
+
+	verbose := flags.Bool(
+		"v",
+		false,
+		"Log progress of each block verification")
+
+	getVersion := flags.Bool(
+		"version",
+		false,
+		"Print version information and exit.")
+
 	// Parse args; omit the first arg which is the command name
-	flags.Parse(os.Args[1:])
+	flags.Parse(args)
+
+	// Print version information if requested
+	if *getVersion {
+		fmt.Printf("keep-block-check %s\n", version)
+		os.Exit(0)
+	}
 
 	config, blobSigningKey, err := loadConfig(*configFile)
 	if err != nil {
@@ -57,18 +84,18 @@ func doMain() error {
 	}
 
 	// get list of block locators to be checked
-	blockLocators, err := getBlockLocators(*locatorFile)
+	blockLocators, err := getBlockLocators(*locatorFile, *prefix)
 	if err != nil {
 		return fmt.Errorf("Error reading block hashes to be checked from file: %s", err.Error())
 	}
 
 	// setup keepclient
-	kc, err := setupKeepClient(config, *keepServicesJSON)
+	kc, blobSignatureTTL, err := setupKeepClient(config, *keepServicesJSON, *blobSignatureTTLFlag)
 	if err != nil {
 		return fmt.Errorf("Error configuring keepclient: %s", err.Error())
 	}
 
-	return performKeepBlockCheck(kc, blobSigningKey, *prefix, blockLocators)
+	return performKeepBlockCheck(kc, blobSignatureTTL, blobSigningKey, blockLocators, *verbose)
 }
 
 type apiConfig struct {
@@ -81,7 +108,7 @@ type apiConfig struct {
 // Load config from given file
 func loadConfig(configFile string) (config apiConfig, blobSigningKey string, err error) {
 	if configFile == "" {
-		err = errors.New("API config file not specified")
+		err = errors.New("Client config file not specified")
 		return
 	}
 
@@ -89,8 +116,6 @@ func loadConfig(configFile string) (config apiConfig, blobSigningKey string, err
 	return
 }
 
-var matchTrue = regexp.MustCompile("^(?i:1|yes|true)$")
-
 // Read config from file
 func readConfigFromFile(filename string) (config apiConfig, blobSigningKey string, err error) {
 	if !strings.Contains(filename, "/") {
@@ -110,20 +135,22 @@ func readConfigFromFile(filename string) (config apiConfig, blobSigningKey strin
 		}
 
 		kv := strings.SplitN(line, "=", 2)
-		key := strings.TrimSpace(kv[0])
-		value := strings.TrimSpace(kv[1])
-
-		switch key {
-		case "ARVADOS_API_TOKEN":
-			config.APIToken = value
-		case "ARVADOS_API_HOST":
-			config.APIHost = value
-		case "ARVADOS_API_HOST_INSECURE":
-			config.APIHostInsecure = matchTrue.MatchString(value)
-		case "ARVADOS_EXTERNAL_CLIENT":
-			config.ExternalClient = matchTrue.MatchString(value)
-		case "ARVADOS_BLOB_SIGNING_KEY":
-			blobSigningKey = value
+		if len(kv) == 2 {
+			key := strings.TrimSpace(kv[0])
+			value := strings.TrimSpace(kv[1])
+
+			switch key {
+			case "ARVADOS_API_TOKEN":
+				config.APIToken = value
+			case "ARVADOS_API_HOST":
+				config.APIHost = value
+			case "ARVADOS_API_HOST_INSECURE":
+				config.APIHostInsecure = arvadosclient.StringBool(value)
+			case "ARVADOS_EXTERNAL_CLIENT":
+				config.ExternalClient = arvadosclient.StringBool(value)
+			case "ARVADOS_BLOB_SIGNING_KEY":
+				blobSigningKey = value
+			}
 		}
 	}
 
@@ -131,7 +158,7 @@ func readConfigFromFile(filename string) (config apiConfig, blobSigningKey strin
 }
 
 // setup keepclient using the config provided
-func setupKeepClient(config apiConfig, keepServicesJSON string) (kc *keepclient.KeepClient, err error) {
+func setupKeepClient(config apiConfig, keepServicesJSON string, blobSignatureTTL time.Duration) (kc *keepclient.KeepClient, ttl time.Duration, err error) {
 	arv := arvadosclient.ArvadosClient{
 		ApiToken:    config.APIToken,
 		ApiServer:   config.APIHost,
@@ -141,7 +168,7 @@ func setupKeepClient(config apiConfig, keepServicesJSON string) (kc *keepclient.
 		External: config.ExternalClient,
 	}
 
-	// if keepServicesJSON is provided, use it to load services; else, use DiscoverKeepServers
+	// If keepServicesJSON is provided, use it instead of service discovery
 	if keepServicesJSON == "" {
 		kc, err = keepclient.MakeKeepClient(&arv)
 		if err != nil {
@@ -155,63 +182,72 @@ func setupKeepClient(config apiConfig, keepServicesJSON string) (kc *keepclient.
 		}
 	}
 
+	// Get if blobSignatureTTL is not provided
+	ttl = blobSignatureTTL
+	if blobSignatureTTL == 0 {
+		value, err := arv.Discovery("blobSignatureTtl")
+		if err == nil {
+			ttl = time.Duration(int(value.(float64))) * time.Second
+		} else {
+			return nil, 0, err
+		}
+	}
+
 	return
 }
 
 // Get list of unique block locators from the given file
-func getBlockLocators(locatorFile string) (locators []string, err error) {
+func getBlockLocators(locatorFile, prefix string) (locators []string, err error) {
 	if locatorFile == "" {
 		err = errors.New("block-hash-file not specified")
 		return
 	}
 
 	content, err := ioutil.ReadFile(locatorFile)
-
 	if err != nil {
 		return
 	}
 
-	locatorMap := make(map[string]string)
-	lines := strings.Split(string(content), "\n")
-	for _, line := range lines {
-		if line == "" {
+	locatorMap := make(map[string]bool)
+	for _, line := range strings.Split(string(content), "\n") {
+		line = strings.TrimSpace(line)
+		if line == "" || !strings.HasPrefix(line, prefix) || locatorMap[line] {
 			continue
 		}
-		trimmedLine := strings.TrimSpace(line)
-		locatorMap[trimmedLine] = trimmedLine
-	}
-
-	for _, locator := range locatorMap {
-		locators = append(locators, locator)
+		locators = append(locators, line)
+		locatorMap[line] = true
 	}
 
 	return
 }
 
 // Get block headers from keep. Log any errors.
-func performKeepBlockCheck(kc *keepclient.KeepClient, blobSigningKey, prefix string, blockLocators []string) error {
-	totalBlocks := 0
+func performKeepBlockCheck(kc *keepclient.KeepClient, blobSignatureTTL time.Duration, blobSigningKey string, blockLocators []string, verbose bool) error {
+	totalBlocks := len(blockLocators)
 	notFoundBlocks := 0
+	current := 0
 	for _, locator := range blockLocators {
-		if !strings.HasPrefix(locator, prefix) {
-			continue
+		current++
+		if verbose {
+			log.Printf("Verifying block %d of %d: %v", current, totalBlocks, locator)
 		}
-
-		totalBlocks++
 		getLocator := locator
 		if blobSigningKey != "" {
 			expiresAt := time.Now().AddDate(0, 0, 1)
-			getLocator = keepclient.SignLocator(locator, kc.Arvados.ApiToken, expiresAt, []byte(blobSigningKey))
+			getLocator = keepclient.SignLocator(locator, kc.Arvados.ApiToken, expiresAt, blobSignatureTTL, []byte(blobSigningKey))
 		}
 
 		_, _, err := kc.Ask(getLocator)
 		if err != nil {
 			notFoundBlocks++
-			log.Printf("Error getting head info for block: %v %v", locator, err)
+			log.Printf("Error verifying block %v: %v", locator, err)
 		}
 	}
+
+	log.Printf("Verify block totals: %d attempts, %d successes, %d errors", totalBlocks, totalBlocks-notFoundBlocks, notFoundBlocks)
+
 	if notFoundBlocks > 0 {
-		return fmt.Errorf("Head information not found for %d out of %d blocks with matching prefix.", notFoundBlocks, totalBlocks)
+		return fmt.Errorf("Block verification failed for %d out of %d blocks with matching prefix.", notFoundBlocks, totalBlocks)
 	}
 
 	return nil