go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/swarming/server/model/utils.go

// Copyright 2023 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package model

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"sort"
	"strconv"
	"strings"

	"github.com/klauspost/compress/zlib"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/gae/service/datastore"

	apipb "go.chromium.org/luci/swarming/proto/api_v2"
)

// checkIsHex returns an error if the string doesn't look like a lowercase hex
// string.
func checkIsHex(s string, minLen int) error {
	if len(s) < minLen {
		return errors.New("too small")
	}
	for _, c := range s {
		if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) {
			return errors.Reason("bad lowercase hex string %q, wrong char %c", s, c).Err()
		}
	}
	return nil
}

// ToJSONProperty serializes a value into a JSON blob property.
//
// Empty maps and lists are stored as nulls.
func ToJSONProperty(val any) (datastore.Property, error) {
	if val == nil {
		return datastore.MkPropertyNI(nil), nil
	}
	blob, err := json.Marshal(val)
	if bytes.Equal(blob, []byte("{}")) || bytes.Equal(blob, []byte("[]")) {
		return datastore.MkPropertyNI(nil), nil
	}
	return datastore.MkPropertyNI(string(blob)), err
}

// FromJSONProperty deserializes a JSON blob property into `val`.
//
// If the property is missing, `val` will be unchanged. Assumes `val` is a list
// or a dict.
//
// Recognizes zlib-compressed properties for compatibility with older entities.
func FromJSONProperty(prop datastore.Property, val any) error {
	propVal, err := prop.Project(datastore.PTBytes)
	if err != nil {
		return err
	}
	blob, _ := propVal.([]byte)
	if len(blob) == 0 {
		return nil
	}

	// This seems to be an uncompressed JSON. Load it as is. Note that zlib
	// compressed data always starts with 0x78 byte (part of the zlib header).
	if blob[0] == '{' || blob[0] == '[' {
		return json.Unmarshal(blob, val)
	}

	// If this doesn't look like JSON, this is likely an older zlib-compressed
	// property. Try to uncompress and load it.
	r, err := zlib.NewReader(bytes.NewBuffer(blob))
	if err != nil {
		return err
	}
	w := bytes.NewBuffer(nil)
	if _, err := io.Copy(w, r); err != nil {
		_ = r.Close()
		return err
	}
	if err := r.Close(); err != nil {
		return err
	}
	return json.Unmarshal(w.Bytes(), val)
}

// LegacyProperty is a placeholder for "recognizing" known legacy properties.
//
// Properties of this type are silently discarded when read (and consequently
// not stored back when written). This is useful for dropping properties that
// were known to exist at some point, but which are no longer used by anything
// at all. If we just ignore them completely, they'll end up in `Extra` maps,
// which we want to avoid (`Extra` is only for truly unexpected properties).
type LegacyProperty struct{}

var _ datastore.PropertyConverter = &LegacyProperty{}

// FromProperty implements datastore.PropertyConverter.
func (*LegacyProperty) FromProperty(p datastore.Property) error {
	return nil
}

// ToProperty implements datastore.PropertyConverter.
func (*LegacyProperty) ToProperty() (datastore.Property, error) {
	return datastore.Property{}, datastore.ErrSkipProperty
}

// SortStringPairs sorts string pairs.
// This was stolen from go.chromium.org/luci/buildbucket/protoutil/tag.go
// and should probably be moved to go.chromium.org/luci/common, but that would
// require a larger refactor, hence the following:
// TODO (crbug.com/1508908): remove this once refactored.
func SortStringPairs(pairs []*apipb.StringPair) {
	sort.Slice(pairs, func(i, j int) bool {
		switch {
		case pairs[i].Key < pairs[j].Key:
			return true
		case pairs[i].Key > pairs[j].Key:
			return false
		default:
			return pairs[i].Value < pairs[j].Value
		}
	})
}

// dimensionsFlatToPb converts a list of k:v pairs into []*apipb.StringListPair.
func dimensionsFlatToPb(flat []string) []*apipb.StringListPair {
	// In the vast majority of cases `flat` is already sorted and we can skip
	// unnecessary maps and resorting. Start with the assumption it is sorted and
	// fallback to a generic implementation if we notice a violation.
	var out []*apipb.StringListPair
	for _, kv := range flat {
		k, v, _ := strings.Cut(kv, ":")
		if len(out) == 0 {
			out = append(out, &apipb.StringListPair{
				Key:   k,
				Value: []string{v},
			})
			continue
		}
		switch prev := out[len(out)-1]; {
		case k == prev.Key:
			switch prevV := prev.Value[len(prev.Value)-1]; {
			case v == prevV:
				// Skip the duplicate.
			case v > prevV:
				prev.Value = append(prev.Value, v)
			default: // v < prevV => the `flat` is not sorted in ascending order
				return dimensionsFlatToPbSlow(flat)
			}
		case k > prev.Key:
			out = append(out, &apipb.StringListPair{
				Key:   k,
				Value: []string{v},
			})
		default: // i.e. k < prev.Key => the `flat` is not sorted in ascending order
			return dimensionsFlatToPbSlow(flat)
		}
	}
	return out
}

// dimensionsFlatToPbSlow is the same as dimensionsFlatToPb, but it doesn't rely
// on `flat` being presorted.
func dimensionsFlatToPbSlow(flat []string) []*apipb.StringListPair {
	sortedCopy := append(make([]string, 0, len(flat)), flat...)
	sort.Strings(sortedCopy)
	return dimensionsFlatToPb(sortedCopy)
}

// MapToStringListPair converts a map[string][]string to []*apipb.StringListPair.
// If keySorting, sorting is applied to the keys.
func MapToStringListPair(p map[string][]string, keySorting bool) []*apipb.StringListPair {
	if len(p) == 0 {
		return nil
	}
	keys := make([]string, 0, len(p))
	for k := range p {
		keys = append(keys, k)
	}
	if keySorting {
		sort.Strings(keys)
	}
	slp := make([]*apipb.StringListPair, len(keys))
	for i, key := range keys {
		slp[i] = &apipb.StringListPair{
			Key:   key,
			Value: p[key],
		}
	}
	return slp
}

// PutMockTaskOutput is a testing util that will create mock TaskOutputChunk datastore entities.
func PutMockTaskOutput(ctx context.Context, reqKey *datastore.Key, numChunks int) {
	toPut := make([]*TaskOutputChunk, numChunks)
	for i := 0; i < numChunks; i++ {
		expectedStr := strings.Repeat(strconv.Itoa(i), ChunkSize)
		var b bytes.Buffer
		w := zlib.NewWriter(&b)
		_, err := w.Write([]byte(expectedStr))
		if err != nil {
			panic(err)
		}
		err = w.Close()
		if err != nil {
			panic(err)
		}
		compressedChunk := b.Bytes()
		toPut[i] = &TaskOutputChunk{
			Key:   TaskOutputChunkKey(ctx, reqKey, int64(i)),
			Chunk: compressedChunk,
		}
	}
	err := datastore.Put(ctx, toPut)
	if err != nil {
		panic(err)
	}
}