irc/message.go

package main

import (
	"bytes"
	"fmt"
	"strings"
)

type Tag struct {
	Key   string
	Value string
}

// Tags is an key-order-preserving, last-insert-wins "set" of Tag
type Tags struct {
	kv   map[string]int // quick lookups for inserts
	tags []Tag
}

func (t *Tags) Insert(new Tag) {
	if t.kv == nil {
		t.kv = map[string]int{}
	}
	i, ok := t.kv[new.Key]
	if ok {
		t.tags[i] = new
	} else {
		t.kv[new.Key] = len(t.tags)
		t.tags = append(t.tags, new)
	}
}

func (t *Tags) GetSlice() []Tag {
	return t.tags
}

type Message struct {
	Tags       []Tag
	Source     string
	Command    command
	Parameters []string
	Raw        string

	parseIndex int
}

func (m *Message) ParseMessage() error {
	m.parseIndex = 0
	m.consumeSpaces()
	if m.consume('@') {
		m.Tags, _ = m.parseTags()
	}
	m.consumeSpaces()
	if m.consume(':') {
		m.Source, _ = m.parseSource()
	}
	m.consumeSpaces()
	m.Command, _ = m.parseCommand()
	m.consumeSpaces()
	m.Parameters, _ = m.parseParameters()
	return nil
}

func (m *Message) parseTags() ([]Tag, error) {
	tags := &Tags{}
	for {
		next := m.findNext("=", ";")
		if next == -1 {
			break
		} else if m.Raw[m.parseIndex+next] == '=' {
			key := m.Raw[m.parseIndex : m.parseIndex+next]
			m.parseIndex += next + 1
			eot := m.findNext(";", " ")
			tags.Insert(Tag{Key: key, Value: unescapeTag(m.Raw[m.parseIndex : m.parseIndex+eot])})
			m.parseIndex += len(m.Raw[m.parseIndex:m.parseIndex+eot]) + 1
		} else if m.Raw[m.parseIndex+next] == ';' {
			key := m.Raw[m.parseIndex : m.parseIndex+next]
			m.parseIndex += next + 1
			tags.Insert(Tag{Key: key, Value: ""})
		} else {
			break
		}
	}
	return tags.GetSlice(), nil
}

func (m *Message) parseSource() (string, error) {
	start := m.parseIndex
	endofparse := strings.Index(m.Raw[m.parseIndex:], " ")
	if endofparse == -1 {
		return "", fmt.Errorf("end of string encountered while parsing tags")
	}
	m.parseIndex += endofparse
	return m.Raw[start:m.parseIndex], nil
}

func (m *Message) parseCommand() (command, error) {
	start := m.parseIndex
	endofparse := strings.Index(m.Raw[m.parseIndex:], " ")
	if endofparse == -1 {
		m.parseIndex = len(m.Raw) // a command with no parameters
	} else {
		m.parseIndex += endofparse
	}
	return command(m.Raw[start:m.parseIndex]), nil
}

func (m *Message) parseParameters() ([]string, error) {
	params := []string{}
	for m.parseIndex <= len(m.Raw) {
		m.consumeSpaces()
		if m.consume(':') { // "run until end of line"
			params = append(params, m.Raw[m.parseIndex:])
			break
		}
		endofparse := strings.Index(m.Raw[m.parseIndex:], " ")
		if endofparse == -1 { // no further params after this one
			p := strings.TrimSuffix(m.Raw[m.parseIndex:], " ")
			if len(p) > 0 { // reject empty trailing params
				params = append(params, p)
			}
			break
		}
		params = append(params, m.Raw[m.parseIndex:m.parseIndex+endofparse])
		m.parseIndex += endofparse
	}
	return params, nil
}

func (m *Message) consumeSpaces() {
	if len(m.Raw) <= m.parseIndex {
		return
	}
	for len(m.Raw) > m.parseIndex && m.Raw[m.parseIndex] == ' ' {
		m.parseIndex++
	}
}

func (m *Message) consume(r byte) bool {
	if len(m.Raw) <= m.parseIndex {
		return false
	}
	if m.Raw[m.parseIndex] == r {
		m.parseIndex++
		return true
	}
	return false
}

func (m *Message) findNext(bs ...string) int {
	if len(bs) == 0 {
		return 0
	}
	r := -1
	for _, b := range bs {
		i := strings.Index(m.Raw[m.parseIndex:], b)
		if r == -1 || (i != -1 && i < r) {
			r = i
		}
	}
	return r
}

func unescapeTag(s string) string {
	sb := strings.Builder{}
	escaping := false
	i := 0
	for i < len(s) {
		if escaping {
			if i+1 > len(s) {
				break
			}
			switch s[i] {
			case ':':
				sb.WriteByte(';')
			case '\\':
				sb.WriteByte('\\')
			case 's':
				sb.WriteByte(' ')
			case 'r':
				sb.WriteByte('\r')
			case 'n':
				sb.WriteByte('\n')
			default:
				sb.WriteByte(s[i])
			}
			i++
			escaping = false
		} else if s[i] == '\\' {
			if !escaping {
				escaping = true
				i++
				continue
			}
		} else {
			sb.WriteByte(s[i])
			i++
		}
	}
	return sb.String()
}

func escapeTag(s string) string {
	sb := strings.Builder{}
	for i := range s {
		switch s[i] {
		case ';':
			sb.Write([]byte{'\\', ':'})
		case ' ':
			sb.Write([]byte{'\\', 's'})
		case '\\':
			sb.Write([]byte{'\\', '\\'})
		case '\n':
			sb.Write([]byte{'\\', 'n'})
		case '\r':
			sb.Write([]byte{'\\', 'r'})
		default:
			sb.WriteByte(s[i])
		}
	}
	return sb.String()
}

func (m *Message) Canonicalize() []byte {
	sb := bytes.NewBuffer([]byte{})
	if m.Tags != nil && len(m.Tags) > 0 {
		sb.WriteByte('@')
		for i := range m.Tags {
			sb.WriteString(m.Tags[i].Key)
			if m.Tags[i].Value != "" {
				sb.WriteByte('=')
				sb.WriteString(escapeTag(m.Tags[i].Value))
			}
			if i+1 != len(m.Tags) {
				sb.WriteByte(';')
			}
		}
		sb.WriteByte(' ')
	}

	if m.Source != "" {
		sb.WriteByte(':')
		sb.WriteString(m.Source)
		sb.WriteByte(' ')
	}

	sb.WriteString(string(m.Command))

	if m.Parameters != nil && len(m.Parameters) > 0 {
		sb.WriteByte(' ')
		for i := range m.Parameters {
			if i+1 == len(m.Parameters) && (strings.Contains(m.Parameters[i], " ") || m.Parameters[i] == "" || m.Parameters[i][0] == ':') {
				sb.WriteByte(':')
			}
			sb.WriteString(m.Parameters[i])
			if i+1 != len(m.Parameters) {
				sb.WriteByte(' ')
			}
		}
	}

	return sb.Bytes()
}