proto.go

     1  // Copyright 2014 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // This file is a simple protocol buffer encoder and decoder.
     6  //
     7  // A protocol message must implement the message interface:
     8  //   decoder() []decoder
     9  //   encode(*buffer)
    10  //
    11  // The decode method returns a slice indexed by field number that gives the
    12  // function to decode that field.
    13  // The encode method encodes its receiver into the given buffer.
    14  //
    15  // The two methods are simple enough to be implemented by hand rather than
    16  // by using a protocol compiler.
    17  //
    18  // See profile.go for examples of messages implementing this interface.
    19  //
    20  // There is no support for groups, message sets, or "has" bits.
    21  
    22  package profile
    23  
    24  import (
    25  	"errors"
    26  	"fmt"
    27  	"slices"
    28  )
    29  
    30  type buffer struct {
    31  	field int
    32  	typ   int
    33  	u64   uint64
    34  	data  []byte
    35  	tmp   [16]byte
    36  }
    37  
    38  type decoder func(*buffer, message) error
    39  
    40  type message interface {
    41  	decoder() []decoder
    42  	encode(*buffer)
    43  }
    44  
    45  func marshal(m message) []byte {
    46  	var b buffer
    47  	m.encode(&b)
    48  	return b.data
    49  }
    50  
    51  func encodeVarint(b *buffer, x uint64) {
    52  	for x >= 128 {
    53  		b.data = append(b.data, byte(x)|0x80)
    54  		x >>= 7
    55  	}
    56  	b.data = append(b.data, byte(x))
    57  }
    58  
    59  func encodeLength(b *buffer, tag int, len int) {
    60  	encodeVarint(b, uint64(tag)<<3|2)
    61  	encodeVarint(b, uint64(len))
    62  }
    63  
    64  func encodeUint64(b *buffer, tag int, x uint64) {
    65  	// append varint to b.data
    66  	encodeVarint(b, uint64(tag)<<3|0)
    67  	encodeVarint(b, x)
    68  }
    69  
    70  func encodeUint64s(b *buffer, tag int, x []uint64) {
    71  	if len(x) > 2 {
    72  		// Use packed encoding
    73  		n1 := len(b.data)
    74  		for _, u := range x {
    75  			encodeVarint(b, u)
    76  		}
    77  		n2 := len(b.data)
    78  		encodeLength(b, tag, n2-n1)
    79  		n3 := len(b.data)
    80  		copy(b.tmp[:], b.data[n2:n3])
    81  		copy(b.data[n1+(n3-n2):], b.data[n1:n2])
    82  		copy(b.data[n1:], b.tmp[:n3-n2])
    83  		return
    84  	}
    85  	for _, u := range x {
    86  		encodeUint64(b, tag, u)
    87  	}
    88  }
    89  
    90  func encodeUint64Opt(b *buffer, tag int, x uint64) {
    91  	if x == 0 {
    92  		return
    93  	}
    94  	encodeUint64(b, tag, x)
    95  }
    96  
    97  func encodeInt64(b *buffer, tag int, x int64) {
    98  	u := uint64(x)
    99  	encodeUint64(b, tag, u)
   100  }
   101  
   102  func encodeInt64Opt(b *buffer, tag int, x int64) {
   103  	if x == 0 {
   104  		return
   105  	}
   106  	encodeInt64(b, tag, x)
   107  }
   108  
   109  func encodeInt64s(b *buffer, tag int, x []int64) {
   110  	if len(x) > 2 {
   111  		// Use packed encoding
   112  		n1 := len(b.data)
   113  		for _, u := range x {
   114  			encodeVarint(b, uint64(u))
   115  		}
   116  		n2 := len(b.data)
   117  		encodeLength(b, tag, n2-n1)
   118  		n3 := len(b.data)
   119  		copy(b.tmp[:], b.data[n2:n3])
   120  		copy(b.data[n1+(n3-n2):], b.data[n1:n2])
   121  		copy(b.data[n1:], b.tmp[:n3-n2])
   122  		return
   123  	}
   124  	for _, u := range x {
   125  		encodeInt64(b, tag, u)
   126  	}
   127  }
   128  
   129  func encodeString(b *buffer, tag int, x string) {
   130  	encodeLength(b, tag, len(x))
   131  	b.data = append(b.data, x...)
   132  }
   133  
   134  func encodeStrings(b *buffer, tag int, x []string) {
   135  	for _, s := range x {
   136  		encodeString(b, tag, s)
   137  	}
   138  }
   139  
   140  func encodeBool(b *buffer, tag int, x bool) {
   141  	if x {
   142  		encodeUint64(b, tag, 1)
   143  	} else {
   144  		encodeUint64(b, tag, 0)
   145  	}
   146  }
   147  
   148  func encodeBoolOpt(b *buffer, tag int, x bool) {
   149  	if !x {
   150  		return
   151  	}
   152  	encodeBool(b, tag, x)
   153  }
   154  
   155  func encodeMessage(b *buffer, tag int, m message) {
   156  	n1 := len(b.data)
   157  	m.encode(b)
   158  	n2 := len(b.data)
   159  	encodeLength(b, tag, n2-n1)
   160  	n3 := len(b.data)
   161  	copy(b.tmp[:], b.data[n2:n3])
   162  	copy(b.data[n1+(n3-n2):], b.data[n1:n2])
   163  	copy(b.data[n1:], b.tmp[:n3-n2])
   164  }
   165  
   166  func unmarshal(data []byte, m message) (err error) {
   167  	b := buffer{data: data, typ: 2}
   168  	return decodeMessage(&b, m)
   169  }
   170  
   171  func le64(p []byte) uint64 {
   172  	return uint64(p[0]) | uint64(p[1])<<8 | uint64(p[2])<<16 | uint64(p[3])<<24 | uint64(p[4])<<32 | uint64(p[5])<<40 | uint64(p[6])<<48 | uint64(p[7])<<56
   173  }
   174  
   175  func le32(p []byte) uint32 {
   176  	return uint32(p[0]) | uint32(p[1])<<8 | uint32(p[2])<<16 | uint32(p[3])<<24
   177  }
   178  
   179  func peekNumVarints(data []byte) (numVarints int) {
   180  	for ; len(data) > 0; numVarints++ {
   181  		var err error
   182  		if _, data, err = decodeVarint(data); err != nil {
   183  			break
   184  		}
   185  	}
   186  	return numVarints
   187  }
   188  
   189  func decodeVarint(data []byte) (uint64, []byte, error) {
   190  	var i int
   191  	var u uint64
   192  	for i = 0; ; i++ {
   193  		if i >= 10 || i >= len(data) {
   194  			return 0, nil, errors.New("bad varint")
   195  		}
   196  		u |= uint64(data[i]&0x7F) << uint(7*i)
   197  		if data[i]&0x80 == 0 {
   198  			return u, data[i+1:], nil
   199  		}
   200  	}
   201  }
   202  
   203  func decodeField(b *buffer, data []byte) ([]byte, error) {
   204  	x, data, err := decodeVarint(data)
   205  	if err != nil {
   206  		return nil, err
   207  	}
   208  	b.field = int(x >> 3)
   209  	b.typ = int(x & 7)
   210  	b.data = nil
   211  	b.u64 = 0
   212  	switch b.typ {
   213  	case 0:
   214  		b.u64, data, err = decodeVarint(data)
   215  		if err != nil {
   216  			return nil, err
   217  		}
   218  	case 1:
   219  		if len(data) < 8 {
   220  			return nil, errors.New("not enough data")
   221  		}
   222  		b.u64 = le64(data[:8])
   223  		data = data[8:]
   224  	case 2:
   225  		var n uint64
   226  		n, data, err = decodeVarint(data)
   227  		if err != nil {
   228  			return nil, err
   229  		}
   230  		if n > uint64(len(data)) {
   231  			return nil, errors.New("too much data")
   232  		}
   233  		b.data = data[:n]
   234  		data = data[n:]
   235  	case 5:
   236  		if len(data) < 4 {
   237  			return nil, errors.New("not enough data")
   238  		}
   239  		b.u64 = uint64(le32(data[:4]))
   240  		data = data[4:]
   241  	default:
   242  		return nil, fmt.Errorf("unknown wire type: %d", b.typ)
   243  	}
   244  
   245  	return data, nil
   246  }
   247  
   248  func checkType(b *buffer, typ int) error {
   249  	if b.typ != typ {
   250  		return errors.New("type mismatch")
   251  	}
   252  	return nil
   253  }
   254  
   255  func decodeMessage(b *buffer, m message) error {
   256  	if err := checkType(b, 2); err != nil {
   257  		return err
   258  	}
   259  	dec := m.decoder()
   260  	data := b.data
   261  	for len(data) > 0 {
   262  		// pull varint field# + type
   263  		var err error
   264  		data, err = decodeField(b, data)
   265  		if err != nil {
   266  			return err
   267  		}
   268  		if b.field >= len(dec) || dec[b.field] == nil {
   269  			continue
   270  		}
   271  		if err := dec[b.field](b, m); err != nil {
   272  			return err
   273  		}
   274  	}
   275  	return nil
   276  }
   277  
   278  func decodeInt64(b *buffer, x *int64) error {
   279  	if err := checkType(b, 0); err != nil {
   280  		return err
   281  	}
   282  	*x = int64(b.u64)
   283  	return nil
   284  }
   285  
   286  func decodeInt64s(b *buffer, x *[]int64) error {
   287  	if b.typ == 2 {
   288  		// Packed encoding
   289  		dataLen := peekNumVarints(b.data)
   290  		*x = slices.Grow(*x, dataLen)
   291  
   292  		data := b.data
   293  		for len(data) > 0 {
   294  			var u uint64
   295  			var err error
   296  
   297  			if u, data, err = decodeVarint(data); err != nil {
   298  				return err
   299  			}
   300  			*x = append(*x, int64(u))
   301  		}
   302  		return nil
   303  	}
   304  	var i int64
   305  	if err := decodeInt64(b, &i); err != nil {
   306  		return err
   307  	}
   308  	*x = append(*x, i)
   309  	return nil
   310  }
   311  
   312  func decodeUint64(b *buffer, x *uint64) error {
   313  	if err := checkType(b, 0); err != nil {
   314  		return err
   315  	}
   316  	*x = b.u64
   317  	return nil
   318  }
   319  
   320  func decodeUint64s(b *buffer, x *[]uint64) error {
   321  	if b.typ == 2 {
   322  		// Packed encoding
   323  		dataLen := peekNumVarints(b.data)
   324  		*x = slices.Grow(*x, dataLen)
   325  
   326  		data := b.data
   327  		for len(data) > 0 {
   328  			var u uint64
   329  			var err error
   330  
   331  			if u, data, err = decodeVarint(data); err != nil {
   332  				return err
   333  			}
   334  			*x = append(*x, u)
   335  		}
   336  		return nil
   337  	}
   338  	var u uint64
   339  	if err := decodeUint64(b, &u); err != nil {
   340  		return err
   341  	}
   342  	*x = append(*x, u)
   343  	return nil
   344  }
   345  
   346  func decodeString(b *buffer, x *string) error {
   347  	if err := checkType(b, 2); err != nil {
   348  		return err
   349  	}
   350  	*x = string(b.data)
   351  	return nil
   352  }
   353  
   354  func decodeStrings(b *buffer, x *[]string) error {
   355  	var s string
   356  	if err := decodeString(b, &s); err != nil {
   357  		return err
   358  	}
   359  	*x = append(*x, s)
   360  	return nil
   361  }
   362  
   363  func decodeBool(b *buffer, x *bool) error {
   364  	if err := checkType(b, 0); err != nil {
   365  		return err
   366  	}
   367  	if int64(b.u64) == 0 {
   368  		*x = false
   369  	} else {
   370  		*x = true
   371  	}
   372  	return nil
   373  }
   374
View as plain text