github.com/nya3jp/tast@v0.0.0-20230601000426-85c8e4d83a9b/src/go.chromium.org/tast/core/internal/bundle/args.go (about)

     1  // Copyright 2018 The ChromiumOS Authors
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  package bundle
     6  
     7  import (
     8  	"encoding/base64"
     9  	"flag"
    10  	"fmt"
    11  	"io"
    12  	"os"
    13  	"path/filepath"
    14  
    15  	"github.com/golang/protobuf/proto"
    16  
    17  	"go.chromium.org/tast/core/errors"
    18  	"go.chromium.org/tast/core/internal/command"
    19  	"go.chromium.org/tast/core/internal/protocol"
    20  )
    21  
    22  type mode int
    23  
    24  const (
    25  	modeRPC mode = iota
    26  	modeDumpTests
    27  	modeRPCTCP
    28  )
    29  
    30  type dumpFormat int
    31  
    32  const (
    33  	dumpFormatLegacyJSON dumpFormat = iota
    34  	dumpFormatProto
    35  )
    36  
    37  type parsedArgs struct {
    38  	mode       mode
    39  	dumpFormat dumpFormat
    40  	// rpctcp mode only
    41  	port      int
    42  	handshake *protocol.HandshakeRequest
    43  }
    44  
    45  // readArgs parses runtime arguments.
    46  // clArgs contains command-line arguments and is typically os.Args[1:].
    47  // The caller is responsible for performing the requested action.
    48  func readArgs(clArgs []string, stderr io.Writer) (*parsedArgs, error) {
    49  	flags := flag.NewFlagSet("", flag.ContinueOnError)
    50  	flags.SetOutput(stderr)
    51  	flags.Usage = func() {
    52  		fmt.Fprintf(stderr, "Usage: %s [flag]...\n\n"+
    53  			"Tast test bundle containing integration tests.\n\n",
    54  			filepath.Base(os.Args[0]))
    55  		flags.PrintDefaults()
    56  	}
    57  
    58  	dump := flags.Bool("dumptests", false, "dump all tests as a JSON-marshaled array of testing.Test structs")
    59  	exportMetadata := flags.Bool("exportmetadata", false, "export all test metadata as a protobuf-marshaled message")
    60  	rpc := flags.Bool("rpc", false, "run gRPC server")
    61  	rpctcp := flags.Bool("rpctcp", false, "run gRPC server listening on TCP. Sample usage:\n"+
    62  		"  cros -rpctcp -port 4444 -handshake [HANDSHAKE_BASE64]")
    63  	port := flags.Int("port", 4444, "port number for gRPC server. Only applicable for rpctcp mode")
    64  	handshakeUsage := "Handshake request for setting up gRPC server. Only applicable for rpctcp mode.\n" +
    65  		"Request should adhere to handshake.proto and be encoded in base64. Example in golang:\n" +
    66  		"  var req *HandshakeRequest = ...\n" +
    67  		"  raw, _ := proto.Marshal(req)\n" +
    68  		"  input = base64.StdEncoding.EncodeToString(raw)"
    69  	handshakeBase64 := flags.String("handshake", "", handshakeUsage)
    70  	if err := flags.Parse(clArgs); err != nil {
    71  		return nil, command.NewStatusErrorf(statusBadArgs, "%v", err)
    72  	}
    73  
    74  	if *dump {
    75  		return &parsedArgs{mode: modeDumpTests, dumpFormat: dumpFormatLegacyJSON}, nil
    76  	}
    77  	if *exportMetadata {
    78  		return &parsedArgs{mode: modeDumpTests, dumpFormat: dumpFormatProto}, nil
    79  	}
    80  	if *rpc {
    81  		return &parsedArgs{mode: modeRPC}, nil
    82  	}
    83  	if *rpctcp {
    84  		var handshakeReq protocol.HandshakeRequest
    85  		if err := decodeBase64Proto(*handshakeBase64, &handshakeReq); err != nil {
    86  			return nil, command.NewStatusErrorf(statusBadArgs, "failed to decode handshake into proto: %v", err)
    87  		}
    88  		return &parsedArgs{
    89  			mode:      modeRPCTCP,
    90  			port:      *port,
    91  			handshake: &handshakeReq,
    92  		}, nil
    93  	}
    94  	flags.Usage()
    95  	return nil, errors.New("no mode flag is set")
    96  }
    97  
    98  // decodeBase64Proto decodes a base64 encoded proto into proto message.
    99  func decodeBase64Proto(protoBase64 string, msg proto.Message) error {
   100  	protoBytes, err := base64.StdEncoding.DecodeString(protoBase64)
   101  	if err != nil {
   102  		return err
   103  	}
   104  	return proto.Unmarshal(protoBytes, msg)
   105  }