go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/logdog/client/butlerlib/streamproto/properties.go (about)

     1  // Copyright 2015 The LUCI Authors.
     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 streamproto
    16  
    17  import (
    18  	"bytes"
    19  	"encoding/json"
    20  	"io"
    21  
    22  	"google.golang.org/protobuf/types/known/timestamppb"
    23  
    24  	"go.chromium.org/luci/common/clock/clockflag"
    25  	"go.chromium.org/luci/common/data/recordio"
    26  	"go.chromium.org/luci/common/errors"
    27  	"go.chromium.org/luci/logdog/api/logpb"
    28  	"go.chromium.org/luci/logdog/common/types"
    29  )
    30  
    31  // Flags is a flag- and JSON-compatible version of logpb.LogStreamDescriptor.
    32  // It is used for stream negotiation protocol and command-line interfaces.
    33  //
    34  // TODO(iannucci) - Change client->butler protocol to just use jsonpb encoding
    35  // of LogStreamDescriptor.
    36  type Flags struct {
    37  	Name        StreamNameFlag `json:"name,omitempty"`
    38  	ContentType string         `json:"contentType,omitempty"`
    39  	Type        StreamType     `json:"type,omitempty"`
    40  	Timestamp   clockflag.Time `json:"timestamp,omitempty"`
    41  	Tags        TagMap         `json:"tags,omitempty"`
    42  }
    43  
    44  // Descriptor converts the Flags to a LogStreamDescriptor.
    45  func (f *Flags) Descriptor() *logpb.LogStreamDescriptor {
    46  	contentType := types.ContentType(f.ContentType)
    47  	if contentType == "" {
    48  		contentType = f.Type.DefaultContentType()
    49  	}
    50  
    51  	var t *timestamppb.Timestamp
    52  	if !f.Timestamp.Time().IsZero() {
    53  		t = timestamppb.New(f.Timestamp.Time())
    54  	}
    55  
    56  	return &logpb.LogStreamDescriptor{
    57  		Name:        string(f.Name),
    58  		ContentType: string(contentType),
    59  		StreamType:  logpb.StreamType(f.Type),
    60  		Timestamp:   t,
    61  		Tags:        f.Tags,
    62  	}
    63  }
    64  
    65  // The maximum size of the initial header in bytes that FromHandshake is willing
    66  // to read.
    67  //
    68  // This must include all fields of the Flags; We're counting on the total
    69  // encoded size of this being less than 1MiB.
    70  const maxFrameSize = 1024 * 1024
    71  
    72  // WriteHandshake writes the butler protocol header handshake on the given
    73  // Writer.
    74  func (f *Flags) WriteHandshake(w io.Writer) error {
    75  	data, err := json.Marshal(f)
    76  	if err != nil {
    77  		return errors.Annotate(err, "marshaling flags").Err()
    78  	}
    79  	if _, err := w.Write(ProtocolFrameHeaderMagic); err != nil {
    80  		return errors.Annotate(err, "writing magic number").Err()
    81  	}
    82  	if _, err := recordio.WriteFrame(w, data); err != nil {
    83  		return errors.Annotate(err, "writing properties").Err()
    84  	}
    85  	return nil
    86  }
    87  
    88  // FromHandshake reads the butler protocol header handshake from the given
    89  // Reader.
    90  func (f *Flags) FromHandshake(r io.Reader) error {
    91  	header := make([]byte, len(ProtocolFrameHeaderMagic))
    92  	_, err := io.ReadFull(r, header)
    93  	if err != nil {
    94  		return errors.Annotate(err, "reading magic number").Err()
    95  	}
    96  	if !bytes.Equal(header, ProtocolFrameHeaderMagic) {
    97  		return errors.Reason(
    98  			"magic number mismatch: got(%q) expected(%q)",
    99  			header, ProtocolFrameHeaderMagic).Err()
   100  	}
   101  
   102  	_, frameReader, err := recordio.NewReader(r, maxFrameSize).ReadFrame()
   103  	if err != nil {
   104  		return errors.Annotate(err, "reading property frame").Err()
   105  	}
   106  
   107  	if err := json.NewDecoder(frameReader).Decode(f); err != nil {
   108  		return errors.Annotate(err, "parsing flag JSON").Err()
   109  	}
   110  
   111  	if frameReader.N > 0 {
   112  		return errors.Reason("handshake had %d bytes of trailing data", frameReader.N).Err()
   113  	}
   114  
   115  	return nil
   116  }