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

     1  // Copyright 2020 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 devserver
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"io"
    11  	"net/http"
    12  	"net/url"
    13  	"os"
    14  
    15  	"github.com/golang/protobuf/ptypes"
    16  	"go.chromium.org/chromiumos/config/go/api/test/tls"
    17  	"go.chromium.org/chromiumos/config/go/api/test/tls/dependencies/longrunning"
    18  	"google.golang.org/grpc"
    19  	"google.golang.org/grpc/codes"
    20  	"google.golang.org/grpc/status"
    21  
    22  	"go.chromium.org/tast/core/errors"
    23  )
    24  
    25  // TLWClient is an implementation of Client to communicate with Test Lab Services wiring API.
    26  type TLWClient struct {
    27  	dutName string
    28  	conn    *grpc.ClientConn
    29  }
    30  
    31  var _ Client = &TLWClient{}
    32  
    33  // NewTLWClient creates a TLWClient.
    34  func NewTLWClient(ctx context.Context, tlwserver, dutName string) (*TLWClient, error) {
    35  	conn, err := grpc.Dial(tlwserver, grpc.WithInsecure())
    36  	if err != nil {
    37  		return nil, errors.Wrapf(err, "failed to establish connection to server: %s", tlwserver)
    38  	}
    39  	return &TLWClient{
    40  		dutName: dutName,
    41  		conn:    conn,
    42  	}, nil
    43  }
    44  
    45  // TearDown closes the gRPC connection to the TLW service.
    46  func (c *TLWClient) TearDown() error {
    47  	return c.conn.Close()
    48  }
    49  
    50  // Stage downloads a file on GCS from storage.googleapis.com using the TLW API.
    51  func (c *TLWClient) Stage(ctx context.Context, gsURL string) (*url.URL, error) {
    52  	// verify GS URL format.
    53  	if _, _, err := ParseGSURL(gsURL); err != nil {
    54  		return nil, errors.Wrapf(err, "failed to parse GS URL: %s", gsURL)
    55  	}
    56  
    57  	req := tls.CacheForDutRequest{Url: gsURL, DutName: c.dutName}
    58  	cl := tls.NewWiringClient(c.conn)
    59  	op, err := cl.CacheForDut(ctx, &req)
    60  	if err != nil {
    61  		st, ok := status.FromError(err)
    62  		if !ok {
    63  			return nil, errors.Wrapf(err, "failed to get status code")
    64  		}
    65  		if st.Code() == codes.NotFound {
    66  			return nil, errors.Wrap(os.ErrNotExist, gsURL)
    67  		}
    68  		return nil, errors.Wrapf(err, "failed to call CacheForDut(%v)", &req)
    69  	}
    70  
    71  	opcli := longrunning.NewOperationsClient(c.conn)
    72  	op, err = opcli.WaitOperation(ctx, &longrunning.WaitOperationRequest{
    73  		Name: op.GetName(),
    74  	})
    75  	if err != nil {
    76  		return nil, errors.Wrap(err, "failed to wait operation")
    77  	}
    78  	if !op.GetDone() {
    79  		return nil, fmt.Errorf("WaitOperation timed out (%v)", op)
    80  	}
    81  
    82  	resp := &tls.CacheForDutResponse{}
    83  	if err := ptypes.UnmarshalAny(op.GetResponse(), resp); err != nil {
    84  		return nil, errors.Wrapf(err, "failed to unmarshal response: %v", resp)
    85  	}
    86  	return url.Parse(resp.Url)
    87  }
    88  
    89  // Open downloads a file on GCS from storage.googleapis.com using the TLW API.
    90  func (c *TLWClient) Open(ctx context.Context, gsURL string) (io.ReadCloser, error) {
    91  	url, err := c.Stage(ctx, gsURL)
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  	httpReq, err := http.NewRequest("GET", url.String(), nil)
    96  	if err != nil {
    97  		return nil, errors.Wrapf(err, "failed to create new HTTP request: %s", url)
    98  	}
    99  	httpReq = httpReq.WithContext(ctx)
   100  
   101  	res, err := defaultHTTPClient.Do(httpReq)
   102  	if err != nil {
   103  		return nil, errors.Wrapf(err, "failed to get from download URL: %s", url)
   104  	}
   105  
   106  	switch res.StatusCode {
   107  	case http.StatusOK:
   108  		return res.Body, nil
   109  	case http.StatusNotFound:
   110  		res.Body.Close()
   111  		return nil, os.ErrNotExist
   112  	default:
   113  		res.Body.Close()
   114  		return nil, fmt.Errorf("got status %d %v", res.StatusCode, httpReq)
   115  	}
   116  }