go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/quota/internal/quotakeys/asi.go (about)

     1  // Copyright 2022 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 quotakeys
    16  
    17  import (
    18  	"encoding/ascii85"
    19  	"strings"
    20  
    21  	"go.chromium.org/luci/common/errors"
    22  )
    23  
    24  const (
    25  	// ASIFieldDelim is used to delimit sections within Application Specific
    26  	// Identifiers (ASIs).
    27  	//
    28  	// NOTE: this is ascii85-safe.
    29  	ASIFieldDelim = "|"
    30  
    31  	// EncodedSectionPrefix is the prefix used for ASI sections which are nominally
    32  	// encoded with ascii85.
    33  	//
    34  	// NOTE: this is ascii85-safe.
    35  	EncodedSectionPrefix     = "{"
    36  	encodedSectionPrefixRune = '{'
    37  
    38  	// EscapedCharacters is the set of characters which are reserved within
    39  	// Application-specific-identifiers (ASIs) and will cause the ASI section to be
    40  	// escaped with ascii85.
    41  	//
    42  	// We also encode ASI sections which start with `EncodedSectionPrefix`.
    43  	EscapedCharacters = QuotaFieldDelim + ASIFieldDelim
    44  )
    45  
    46  func extendBuffer(buf []byte, n int) []byte {
    47  	if cap(buf)-len(buf) < n {
    48  		buf = append(make([]byte, 0, len(buf)+n), buf...)
    49  	}
    50  	return buf[:n]
    51  }
    52  
    53  // AssembleASI will return an ASI with the given sections.
    54  //
    55  // Sections are assembled with a "|" separator verbatim, unless the section
    56  // contains a "|", "~" or begins with "{". In this case the section will be
    57  // encoded with ascii85 and inserted to the final string with a "{" prefix
    58  // character.
    59  func AssembleASI(sections ...string) string {
    60  	encodedSections := make([]string, len(sections))
    61  	var buf []byte
    62  
    63  	for i, section := range sections {
    64  		if strings.HasPrefix(section, EncodedSectionPrefix) || strings.ContainsAny(section, EscapedCharacters) {
    65  			buf = extendBuffer(buf, ascii85.MaxEncodedLen(len(section))+1)
    66  			buf[0] = encodedSectionPrefixRune
    67  			smallBuf := buf[1:]
    68  			encodedSections[i] = string(buf[:ascii85.Encode(smallBuf, []byte(section))+1])
    69  		} else {
    70  			encodedSections[i] = section
    71  		}
    72  	}
    73  	return strings.Join(encodedSections, ASIFieldDelim)
    74  }
    75  
    76  // DecodeASI will return the sections within an ASI, decoding any which appear
    77  // to be ascii85-encoded.
    78  //
    79  // If a section has the ascii85 prefix, but doesn't correctly decode, this
    80  // returns an error.
    81  func DecodeASI(asi string) ([]string, error) {
    82  	if asi == "" {
    83  		return nil, nil
    84  	}
    85  	var buf []byte
    86  	sections := strings.Split(asi, ASIFieldDelim)
    87  	for i, section := range sections {
    88  		if strings.HasPrefix(section, EncodedSectionPrefix) {
    89  			buf = extendBuffer(buf, len(section)-1)
    90  			ndst, _, err := ascii85.Decode(buf, []byte(section[1:]), true)
    91  			if err != nil {
    92  				return nil, errors.Annotate(err, "DecodeASI: section[%d]", i).Err()
    93  			}
    94  			sections[i] = string(buf[:ndst])
    95  		}
    96  	}
    97  	return sections, nil
    98  }