apcupsd_exporter/vendor/github.com/mdlayher/apcupsd/nis.go

83 lines
2.0 KiB
Go

package apcupsd
import (
"encoding/binary"
"errors"
"io"
"math"
"sync"
)
var _ io.ReadWriteCloser = &nisReadWriteCloser{}
// newNISReadWriteCloser wraps an io.ReadWriteCloser.
func newNISReadWriteCloser(rwc io.ReadWriteCloser) *nisReadWriteCloser {
return &nisReadWriteCloser{
rwc: rwc,
lenb: make([]byte, 2),
}
}
// An nisReadWriteCloser wraps an io.ReadWriteCloser with one that can encode
// and decode messages using the NIS's protocol.
type nisReadWriteCloser struct {
mu sync.Mutex
rwc io.ReadWriteCloser
lenb []byte
}
// Read reads messages from the NIS using its protocol:
// - 2 bytes: length of next message
// - N bytes: data
func (rwc *nisReadWriteCloser) Read(b []byte) (int, error) {
rwc.mu.Lock()
defer rwc.mu.Unlock()
// Read two byte length of next data
if _, err := io.ReadFull(rwc.rwc, rwc.lenb); err != nil {
return 0, err
}
// When no more data returned from server, return io.EOF
length := binary.BigEndian.Uint16(rwc.lenb)
if length == 0 {
return 0, io.EOF
}
return io.ReadFull(rwc.rwc, b[:length])
}
var (
// errBufferTooLarge indicates that nisReadWriteCloser.Write was passed a
// buffer that is too large to send to the NIS.
errBufferTooLarge = errors.New("buffer too large; must be size of uint16 or less")
)
// Write writes messages to the NIS using its protocol by prepending each
// message with its 2 byte length.
func (rwc *nisReadWriteCloser) Write(b []byte) (int, error) {
// Cannot write more than math.MaxUint16 bytes
if len(b) > math.MaxUint16 {
return 0, errBufferTooLarge
}
rwc.mu.Lock()
defer rwc.mu.Unlock()
// Two byte length of data
binary.BigEndian.PutUint16(rwc.lenb, uint16(len(b)))
// Send data and indicate the length of the body to caller
n, err := rwc.rwc.Write(append(rwc.lenb, b...))
n -= len(rwc.lenb)
return n, err
}
// Close closes the underlying io.ReadWriteCloser.
func (rwc *nisReadWriteCloser) Close() error {
rwc.mu.Lock()
defer rwc.mu.Unlock()
return rwc.rwc.Close()
}