github.com/egonelbre/exp@v0.0.0-20240430123955-ed1d3aa93911/ber/forkable.go (about)

     1  package ber
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  )
     7  
     8  type forkableWriter struct {
     9  	*bytes.Buffer
    10  	pre, post *forkableWriter
    11  }
    12  
    13  func newForkableWriter() *forkableWriter {
    14  	return &forkableWriter{new(bytes.Buffer), nil, nil}
    15  }
    16  func (f *forkableWriter) fork() (pre, post *forkableWriter) {
    17  	if f.pre != nil || f.post != nil {
    18  		panic("have already forked")
    19  	}
    20  	f.pre = newForkableWriter()
    21  	f.post = newForkableWriter()
    22  	return f.pre, f.post
    23  }
    24  func (f *forkableWriter) Len() (l int) {
    25  	l += f.Buffer.Len()
    26  	if f.pre != nil {
    27  		l += f.pre.Len()
    28  	}
    29  	if f.post != nil {
    30  		l += f.post.Len()
    31  	}
    32  	return
    33  }
    34  
    35  func (f *forkableWriter) writeTo(out io.Writer) (n int, err error) {
    36  	n, err = out.Write(f.Bytes())
    37  	if err != nil {
    38  		return
    39  	}
    40  
    41  	var nn int
    42  
    43  	if f.pre != nil {
    44  		nn, err = f.pre.writeTo(out)
    45  		n += nn
    46  		if err != nil {
    47  			return
    48  		}
    49  	}
    50  
    51  	if f.post != nil {
    52  		nn, err = f.post.writeTo(out)
    53  		n += nn
    54  	}
    55  	return
    56  }
    57  
    58  func marshalHeader(out *forkableWriter, h header) (err error) {
    59  	b := uint8(h.class) << 6
    60  	if h.constructed {
    61  		b |= 0x20
    62  	}
    63  	if h.tag >= 31 {
    64  		b |= 0x1f
    65  		err = out.WriteByte(b)
    66  		if err != nil {
    67  			return
    68  		}
    69  		err = marshalBase128Int(out, int64(h.tag))
    70  		if err != nil {
    71  			return
    72  		}
    73  	} else {
    74  		b |= uint8(h.tag)
    75  		err = out.WriteByte(b)
    76  		if err != nil {
    77  			return
    78  		}
    79  	}
    80  	if h.length >= 128 {
    81  		l := lengthLength(h.length)
    82  		err = out.WriteByte(0x80 | byte(l))
    83  		if err != nil {
    84  			return
    85  		}
    86  		err = marshalLength(out, h.length)
    87  		if err != nil {
    88  			return
    89  		}
    90  	} else {
    91  		err = out.WriteByte(byte(h.length))
    92  		if err != nil {
    93  			return
    94  		}
    95  	}
    96  	return nil
    97  }
    98  
    99  func marshalBase128Int(out *forkableWriter, n int64) (err error) {
   100  	if n == 0 {
   101  		err = out.WriteByte(0)
   102  		return
   103  	}
   104  
   105  	l := 0
   106  	for i := n; i > 0; i >>= 7 {
   107  		l++
   108  	}
   109  
   110  	for i := l - 1; i >= 0; i-- {
   111  		o := byte(n >> uint(i*7))
   112  		o &= 0x7f
   113  		if i != 0 {
   114  			o |= 0x80
   115  		}
   116  		err = out.WriteByte(o)
   117  		if err != nil {
   118  			return
   119  		}
   120  	}
   121  
   122  	return nil
   123  }
   124  
   125  func marshalLength(out *forkableWriter, i int) (err error) {
   126  	n := lengthLength(i)
   127  
   128  	for ; n > 0; n-- {
   129  		err = out.WriteByte(byte(i >> uint((n-1)*8)))
   130  		if err != nil {
   131  			return
   132  		}
   133  	}
   134  
   135  	return nil
   136  }
   137  
   138  func lengthLength(i int) (numBytes int) {
   139  	numBytes = 1
   140  	for i > 255 {
   141  		numBytes++
   142  		i >>= 8
   143  	}
   144  	return
   145  }