irc/message.go

package main

import (
	"bytes"
	"strings"
)

type Tag struct {
	Key   string
	Value string
}

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

// Insert adds the given tag to the set, overriding any existing tag with that Key
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)
	}
}

// GetSlice returns a simple slice of tags from the set
func (t *Tags) GetSlice() []Tag {
	return t.tags
}

// A Message represents a parsed set of Tags/Source/Command/Parameters from a Raw string.
//
// Tags/Source/Command/Parameters are all set from the Raw after ParseMessage() is called.
// Calling ParseMessage multiple times (without changing Raw) should be a no-op.
//
// Canonicalize() returns Tags/Source/Command/Parameters transformed back into an RFC-compliant []byte.
//
// Reusing a message struct to parse multiple Raw messages is possible, but frowned upon.
//
// Tag values are automatically escaped during parsing.
type Message struct {
	Tags       []Tag    // An ordered slice of Key=Value pairs
	Source     string   // An optional string denoting the source of the message
	Command    command  // A command/numeric for the message
	Parameters []string // Zero or more parameters for the Command.
	Raw        string

	parseIndex int
}

// ParseMessage will attempt to parse whatever it's given as an IRC Message,
// and populate Tags/Source/Command/Parameters with the results.
// Garbage in, garbage out, though. The parser tries its best, but doesn't aggressively
// validate messages.
//
// Tag values are automatically escaped during parsing, as needed.
func (m *Message) ParseMessage() {
	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()
}

func (m *Message) parseTags() []Tag {
	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(";", " ")
			if eot == -1 { // ran out of message, probably
				tags.Insert(Tag{Key: key, Value: unescapeTag(m.Raw[m.parseIndex:])})
				m.parseIndex = len(m.Raw)
				break
			}
			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: ""})
		}
	}
	return tags.GetSlice()
}

func (m *Message) parseSource() string {
	start := m.parseIndex
	endofparse := strings.Index(m.Raw[m.parseIndex:], " ")
	if endofparse == -1 {
		m.parseIndex = len(m.Raw) // out of message! which is weird.
	} else {
		m.parseIndex += endofparse
	}
	return m.Raw[start:m.parseIndex]
}

func (m *Message) parseCommand() command {
	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])
}

func (m *Message) parseParameters() []string {
	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
			break
		}
		params = append(params, m.Raw[m.parseIndex:m.parseIndex+endofparse])
		m.parseIndex += endofparse
	}
	return params
}

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 {
	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 {
			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()
}