github.com/bartle-stripe/trillian@v1.2.1/cmd/updatetree/main.go (about)

     1  // Copyright 2018 Google Inc. All Rights Reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package main contains the implementation and entry point for the updatetree
    16  // command.
    17  //
    18  // Example usage:
    19  // $ ./updatetree --admin_server=host:port --tree_id=123456789 --tree_state=FROZEN
    20  //
    21  // The output is minimal to allow for easy usage in automated scripts.
    22  package main
    23  
    24  import (
    25  	"context"
    26  	"errors"
    27  	"flag"
    28  	"fmt"
    29  	"time"
    30  
    31  	"github.com/golang/glog"
    32  	"github.com/golang/protobuf/proto"
    33  	"github.com/google/trillian"
    34  	"github.com/google/trillian/client/rpcflags"
    35  	"google.golang.org/genproto/protobuf/field_mask"
    36  	"google.golang.org/grpc"
    37  	"google.golang.org/grpc/codes"
    38  	"google.golang.org/grpc/status"
    39  )
    40  
    41  var (
    42  	adminServerAddr = flag.String("admin_server", "", "Address of the gRPC Trillian Admin Server (host:port)")
    43  	rpcDeadline     = flag.Duration("rpc_deadline", time.Second*10, "Deadline for RPC requests")
    44  	treeID          = flag.Int64("tree_id", 0, "The ID of the tree to be set updated")
    45  	treeState       = flag.String("tree_state", "", "If set the tree state will be updated")
    46  )
    47  
    48  // TODO(Martin2112): Pass everything needed into this and don't refer to flags.
    49  func updateTree(ctx context.Context) (*trillian.Tree, error) {
    50  	if *adminServerAddr == "" {
    51  		return nil, errors.New("empty --admin_server, please provide the Admin server host:port")
    52  	}
    53  
    54  	m := proto.EnumValueMap("trillian.TreeState")
    55  	if m == nil {
    56  		return nil, fmt.Errorf("can't find enum value map for states")
    57  	}
    58  	newState, ok := m[*treeState]
    59  	if !ok {
    60  		return nil, fmt.Errorf("invalid tree state: %v", *treeState)
    61  	}
    62  
    63  	// We only want to update the state of the tree, which means we need a field
    64  	// mask on the request.
    65  	treeStateMask := &field_mask.FieldMask{
    66  		Paths: []string{"tree_state"},
    67  	}
    68  
    69  	req := &trillian.UpdateTreeRequest{
    70  		Tree: &trillian.Tree{
    71  			TreeId:    *treeID,
    72  			TreeState: trillian.TreeState(newState),
    73  		},
    74  		UpdateMask: treeStateMask,
    75  	}
    76  
    77  	dialOpts, err := rpcflags.NewClientDialOptionsFromFlags()
    78  	if err != nil {
    79  		return nil, fmt.Errorf("failed to determine dial options: %v", err)
    80  	}
    81  
    82  	conn, err := grpc.Dial(*adminServerAddr, dialOpts...)
    83  	if err != nil {
    84  		return nil, fmt.Errorf("failed to dial %v: %v", *adminServerAddr, err)
    85  	}
    86  	defer conn.Close()
    87  
    88  	client := trillian.NewTrillianAdminClient(conn)
    89  	for {
    90  		tree, err := client.UpdateTree(ctx, req)
    91  		if err == nil {
    92  			return tree, nil
    93  		}
    94  		if s, ok := status.FromError(err); ok && s.Code() == codes.Unavailable {
    95  			glog.Errorf("Admin server unavailable, trying again: %v", err)
    96  			time.Sleep(100 * time.Millisecond)
    97  			continue
    98  		}
    99  		return nil, fmt.Errorf("failed to UpdateTree(%+v): %T %v", req, err, err)
   100  	}
   101  }
   102  
   103  func main() {
   104  	flag.Parse()
   105  	defer glog.Flush()
   106  
   107  	ctx, cancel := context.WithTimeout(context.Background(), *rpcDeadline)
   108  	defer cancel()
   109  	tree, err := updateTree(ctx)
   110  	if err != nil {
   111  		glog.Exitf("Failed to update tree: %v", err)
   112  	}
   113  
   114  	// DO NOT change the output format, scripts are meant to depend on it.
   115  	// If you really want to change it, provide an output_format flag and
   116  	// keep the default as-is.
   117  	fmt.Println(tree.TreeState)
   118  }