lg/internal/socket/socket.go

package socket

// taken from, adapted by me
// https://github.com/StatCan/go-birdc/blob/master/socket/socket.go

import (
	"bufio"
	"errors"
	"fmt"
	"net"
	"regexp"
	"strings"
	"time"
)

var replyCodeExpr *regexp.Regexp
var contentCodeExpr *regexp.Regexp
var codeExpr *regexp.Regexp

func init() {
	// https://gitlab.nic.cz/labs/bird/-/blob/master/doc/reply_codes
	replyCodeExpr = regexp.MustCompile(`(?m)^([089][0-9]{3})`)
	contentCodeExpr = regexp.MustCompile(`(?m)^([12][0-9]{3})`)
	codeExpr = regexp.MustCompile(`(?m)^([0-9]{4})`)
}

// BirdSocket represents a socket connection to bird daemon
//
// `path` is the unix socket path
//
// `bufferSize` is the size of the read buffer
//
// `conn` is the actual socket connection
type BirdSocket struct {
	path       string
	bufferSize int
	conn       net.Conn
	restricted bool
}

// NewBirdSocket creates a new BirdSocket
//
// `path` is the unix socket path
//
// `bufferSize` is the size of the read buffer
func NewBirdSocket(path string, bufferSize int, restricted bool) *BirdSocket {
	return &BirdSocket{path: path, bufferSize: bufferSize, restricted: restricted}
}

// Connect opens the unix socket connection
func (s *BirdSocket) Connect() (err error) {
	s.conn, err = net.Dial("unix", s.path)
	if err != nil {
		return
	}

	// throw away the welcome message
	buf := make([]byte, s.bufferSize)
	_, err = s.conn.Read(buf[:])

	_, resp_code, err := s.Send("restrict")
	if err != nil {
		return err
	}
	if string(resp_code) != "0016" { // access restricted
		return errors.New("failed to restrict bird session")
	}

	return
}

// Close closes the unix socket connection
//
// **NOTE** it is important to close the socket connection since,
// in the event that there is still data waiting to be transmitted over the connection,
// close will try to complete the transmission properly
func (s *BirdSocket) Close() {
	if s.conn != nil {
		s.conn.Close()
	}
}

func (s *BirdSocket) Send(data string) (resp []byte, replyCode []byte, err error) {
	err = s.write([]byte(data))
	if err != nil {
		return
	}
	return s.read()
}

func (s *BirdSocket) write(data []byte) (err error) {
	if s.conn == nil {
		err = fmt.Errorf("unable to write to socket, connection is not open. Ensure that s.Connect() is called first. %s", s.conn)
		return
	}

	err = s.conn.SetWriteDeadline(time.Now().Add(30 * time.Second))
	if err != nil {
		return
	}
	_, err = s.conn.Write([]byte(strings.Trim(string(data), "\n") + "\n"))
	return
}

// read reads data from the socket connection and captures the replyCode that indicated the end of transmissions
// and any associated response or error
//
// For more information on reply codes, see https://gitlab.nic.cz/labs/bird/-/blob/master/doc/reply_codes
//
// `resp` is the raw response from BIRD
//
// `replyCode` is the reply code associated with the response
func (s *BirdSocket) read() (resp []byte, code []byte, err error) {
	err = s.conn.SetReadDeadline(time.Now().Add(30 * time.Second))
	if err != nil {
		return
	}

	resp = make([]byte, 0)
	code = []byte("7000")

	reader := bufio.NewReader(s.conn)

	for !containsReplyCode(code) {
		for {
			line, err_inside := reader.ReadBytes('\n')
			if err_inside != nil {
				return nil, nil, err_inside
			}

			linecode := line[0:4]

			isCode := containsCode(linecode)

			if isCode {
				code = linecode
			}

			// ' ' as 5th char means end of output, otherwise will continue on other line, do not break loop
			if isCode && containsReplyCode(code) {
				// do not append the last line as it will be another '\n'
				resp = append(resp, line[5:len(line)-1]...)
				break
			} else if isCode {
				resp = append(resp, line[5:]...)
			} else if containsContentCode(code) {
				resp = append(resp, line[1:]...)
			} else {
				fallback := fmt.Appendf([]byte("<<<unparsable_string("), string(line), ")>>>")
				resp = append(resp, fallback...)
				err = errors.New("unparsable string")
				break
			}
		}
		if err != nil {
			return
		}
	}

	return
}

func containsCode(b []byte) bool {
	return codeExpr.Match(b)
}

func containsContentCode(b []byte) bool {
	return contentCodeExpr.Match(b)
}

func containsReplyCode(b []byte) bool {
	return replyCodeExpr.Match(b)
}