1// Copyright 2018 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5package osp
6
7import (
8	"encoding/binary"
9	"errors"
10	"io"
11)
12
13const (
14	maxInt6  = 1<<5 - 1
15	maxInt14 = 1<<13 - 1
16	maxInt30 = 1<<29 - 1
17	maxInt62 = 1<<63 - 1
18)
19
20// https://tools.ietf.org/html/draft-ietf-quic-transport-16#section-16
21func ReadVaruint(r io.Reader) (uint64, error) {
22	b0, err := readByte(r)
23	if err != nil {
24		return 0, err
25	}
26	var a [8]byte
27	b := a[:]
28	e := binary.BigEndian
29	b[0] = last6Bits(b0)
30	switch first2Bits(b0) {
31	case 0:
32		return uint64(b[0]), nil
33	case 1:
34		err = readBytes(r, b[1:2])
35		return uint64(e.Uint16(b)), err
36	case 2:
37		err = readBytes(r, b[1:4])
38		return uint64(e.Uint32(b)), err
39	case 3:
40		err = readBytes(r, b[1:8])
41		return uint64(e.Uint64(b)), err
42	}
43	return 0, nil
44}
45
46// https://tools.ietf.org/html/draft-ietf-quic-transport-16#section-16
47func WriteVaruint(v uint64, w io.Writer) error {
48	var a [8]byte
49	b := a[:]
50	e := binary.BigEndian
51	if v <= maxInt6 {
52		b[0] = byte(v)
53		setFirst2Bits(0, b)
54		return writeBytes(b[:1], w)
55	} else if v <= maxInt14 {
56		e.PutUint16(b, uint16(v))
57		setFirst2Bits(1, b)
58		return writeBytes(b[:2], w)
59	} else if v <= maxInt30 {
60		e.PutUint32(b, uint32(v))
61		setFirst2Bits(2, b)
62		return writeBytes(b[:4], w)
63	} else if v <= maxInt62 {
64		e.PutUint64(b, v)
65		setFirst2Bits(3, b)
66		return writeBytes(b[:8], w)
67	}
68	return errors.New("Too big")
69}
70
71func first2Bits(b byte) byte {
72	return b >> 6
73}
74
75func last6Bits(b byte) byte {
76	return b & 0x3f // 0b00111111
77}
78
79func setFirst2Bits(first byte, b []byte) {
80	b[0] = (first << 6) | b[0]
81}
82
83func readByte(r io.Reader) (byte, error) {
84	var b [1]byte
85	err := readBytes(r, b[:])
86	return b[0], err
87}
88
89// Read len(b) bytes from r into b
90// If you hit an error (of the end), propogate the error from r.Read
91func readBytes(r io.Reader, b []byte) error {
92	start := 0
93	end := len(b)
94	for start < end {
95		n, err := r.Read(b[start:end])
96		if err != nil {
97			return err
98		}
99		start += n
100	}
101	return nil
102}
103
104// Write len(b) bytes from b to w
105// If you hit an error, propogate the error from w.Write
106func writeBytes(b []byte, w io.Writer) error {
107	start := 0
108	end := len(b)
109	for start < end {
110		n, err := w.Write(b[start:end])
111		if err != nil {
112			return err
113		}
114		start += n
115	}
116	return nil
117}
118