github.com/google/osv-scalibr@v0.4.1/binary/proto/io.go (about)

     1  // Copyright 2025 Google LLC
     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 proto
    16  
    17  import (
    18  	"compress/gzip"
    19  	"errors"
    20  	"os"
    21  	"path/filepath"
    22  	"strings"
    23  
    24  	"github.com/google/osv-scalibr/log"
    25  	"google.golang.org/protobuf/encoding/prototext"
    26  	"google.golang.org/protobuf/proto"
    27  )
    28  
    29  // fileType represents the type of a proto result file.
    30  type fileType struct {
    31  	isGZipped  bool
    32  	isBinProto bool
    33  }
    34  
    35  // typeForPath returns the proto type of a path, or an error if the path is not a valid proto file.
    36  func typeForPath(filePath string) (*fileType, error) {
    37  	ext := filepath.Ext(filePath)
    38  	if ext == "" {
    39  		return nil, errors.New("invalid filename: Doesn't have an extension")
    40  	}
    41  
    42  	isGZipped := false
    43  	if ext == ".gz" {
    44  		isGZipped = true
    45  		ext = filepath.Ext(strings.TrimSuffix(filePath, ext))
    46  		if ext == "" {
    47  			return nil, errors.New("invalid filename: Gzipped file doesn't have an extension")
    48  		}
    49  	}
    50  
    51  	var isBinProto bool
    52  	switch ext {
    53  	case ".binproto":
    54  		isBinProto = true
    55  	case ".textproto":
    56  		isBinProto = false
    57  	default:
    58  		return nil, errors.New("invalid filename: not a .textproto or .binproto")
    59  	}
    60  
    61  	return &fileType{isGZipped: isGZipped, isBinProto: isBinProto}, nil
    62  }
    63  
    64  // ValidExtension returns an error if the file extension is not a proto file.
    65  func ValidExtension(path string) error {
    66  	_, err := typeForPath(path)
    67  	return err
    68  }
    69  
    70  // Write writes a proto message to a .textproto or .binproto file, based on the file extension.
    71  // If the file name additionally has the .gz suffix, it's zipped before writing.
    72  func Write(filePath string, outputProto proto.Message) error {
    73  	ft, err := typeForPath(filePath)
    74  	if err != nil {
    75  		return err
    76  	}
    77  	return write(filePath, outputProto, ft)
    78  }
    79  
    80  // WriteWithFormat writes a proto message to a .textproto or .binproto file, based
    81  // on the value of the format parameter ("textproto" or "binproto")
    82  func WriteWithFormat(filePath string, outputProto proto.Message, format string) error {
    83  	ft := &fileType{isGZipped: false, isBinProto: format == "binproto"}
    84  	return write(filePath, outputProto, ft)
    85  }
    86  
    87  func write(filePath string, outputProto proto.Message, ft *fileType) error {
    88  	var p []byte
    89  	var err error
    90  	if ft.isBinProto {
    91  		if p, err = proto.Marshal(outputProto); err != nil {
    92  			return err
    93  		}
    94  	} else {
    95  		opts := prototext.MarshalOptions{Multiline: true}
    96  		if p, err = (opts.Marshal(outputProto)); err != nil {
    97  			return err
    98  		}
    99  	}
   100  
   101  	log.Infof("Marshaled result proto has %d bytes", len(p))
   102  
   103  	f, err := os.Create(filePath)
   104  	if err != nil {
   105  		return err
   106  	}
   107  	defer f.Close()
   108  	if ft.isGZipped {
   109  		writer := gzip.NewWriter(f)
   110  		if _, err := writer.Write(p); err != nil {
   111  			return err
   112  		}
   113  		if err := writer.Close(); err != nil {
   114  			return err
   115  		}
   116  	} else if _, err := f.Write(p); err != nil {
   117  		return err
   118  	}
   119  	return nil
   120  }