dndmusicbot/discord/discord.go

package discord

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/gorilla/websocket"
	"github.com/kataras/go-events"
	"github.com/tidwall/gjson"
)

const (
	GATEWAYBOTURL = "https://discord.com/api/gateway/bot"
)

const (
	PermissionVoiceConnect = 0x0000000000100000
	PermissionVoiceSpeak   = 0x0000000000200000
	PermissionSendMessages = 0x0000000000000800
)

type Connection struct {
	events        events.EventEmmiter
	httpClient    *http.Client
	dialer        *websocket.Dialer
	conn          *websocket.Conn
	token         string
	seq           int
	ready         chan bool
	vstateUpdate  chan bool
	vserverUpdate chan bool
}

type VoiceConnection struct {
	d    *Connection
	conn *websocket.Conn
	sid  string
	e    string
}

type VoiceStateUpdate struct {
	Op   int `json:"op"`
	Data struct {
		Guild   string `json:"guild_id"`
		Channel string `json:"channel_id"`
		Mute    bool   `json:"self_mute"`
		Deaf    bool   `json:"self_dead"`
	} `json:"d"`
}

type Identify struct {
	Op int `json:"op"`

	Data struct {
		Token   string `json:"token"`
		Intents int    `json:"intents"`

		Properties struct {
			OS      string `json:"os"`
			Browser string `json:"browser"`
			Device  string `json:"device"`
		} `json:"properties"`
	} `json:"d"`
}

func NewConnection(token string) (*Connection, error) {
	conn := new(Connection)

	conn.token = token

	conn.ready = make(chan bool)
	conn.vserverUpdate = make(chan bool)
	conn.vstateUpdate = make(chan bool)

	conn.events = events.New()
	conn.events.On("READY", conn.eventready)

	conn.dialer = &websocket.Dialer{
		HandshakeTimeout: 45 * time.Second,
	}

	conn.httpClient = &http.Client{}
	gw, err := conn.GetGateway()
	if err != nil {
		return nil, err
	}

	conn.conn, _, err = conn.dialer.Dial(gw, nil)
	if err != nil {
		return nil, err
	}

	go func() {
		for {
			_, b, err := conn.conn.ReadMessage()
			if err != nil {
				log.Println(err)
				break
			}

			go conn.handleGatewayEvent(b)
		}
	}()

	// Identify
	id := new(Identify)
	id.Op = 2
	id.Data.Token = token
	id.Data.Intents = PermissionVoiceConnect | PermissionVoiceSpeak | PermissionSendMessages
	id.Data.Properties.OS = "linux"
	id.Data.Properties.Browser = "dndmusicbot/0.0.1"
	id.Data.Properties.Device = "dndmusicbot"

	json.NewEncoder(os.Stdout).Encode(id)

	err = conn.conn.WriteJSON(id)
	if err != nil {
		conn.conn.Close()
		return nil, err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// Wait for Ready!
	select {
	case <-ctx.Done():
		return nil, err
	case <-conn.ready:
	}

	return conn, nil
}

func (c *Connection) NewVoiceConnection(guild string, channel string) (*VoiceConnection, error) {
	vc := new(VoiceConnection)
	vc.d = c

	c.events.On("VOICE_STATE_UPDATE", vc.voiceStateUpdate)
	c.events.On("VOICE_SERVER_UDPATE", vc.voiceServerUpdate)

	msg := new(VoiceStateUpdate)
	msg.Op = 4
	msg.Data.Guild = guild
	msg.Data.Channel = channel
	msg.Data.Mute = false
	msg.Data.Deaf = true

	c.conn.WriteJSON(msg)

	var state, server bool

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// We have to wait for both a vserverupdate and a vstateupdate before we continue.
	for !state && !server {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-c.vserverUpdate:
			server = true
		case <-c.vstateUpdate:
			state = true
		}
	}

	// We should have gotten all the information by now.. continue..

	return vc, nil
}

func (vc *VoiceConnection) voiceStateUpdate(payload ...interface{}) {
	vc.d.vstateUpdate <- true
}

func (vc *VoiceConnection) voiceServerUpdate(payload ...interface{}) {
	vc.d.vserverUpdate <- true
}

// Ready event. Contains alot of information regarding our session.
func (c *Connection) eventready(payload ...interface{}) {
	c.ready <- true
}

func (c *Connection) handleGatewayEvent(js []byte) error {
	eventdata := gjson.GetManyBytes(js, "op", "d", "s", "t")

	op := eventdata[0]
	data := eventdata[1]
	seq := eventdata[2]
	eventname := eventdata[3]

	if seq.Exists() {
		c.seq = int(seq.Int())
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	c.conn.SetCloseHandler(func(code int, text string) error {
		cancel()
		return nil
	})

	switch op.Int() {
	case 0:
		fmt.Printf("Event: %s Data: %+v\n", eventname.String(), data)
		c.events.Emit(events.EventName(eventname.String()), data)
	case 10: // Hello
		hb_interval, err := time.ParseDuration(fmt.Sprintf("%dms", data.Get("heartbeat_interval").Int()))
		if err != nil {
			return err
		}

		log.Printf("We Received a Hello! Starting heartbeat on %.2fs interval\n", hb_interval.Seconds())

		hb_ticker := time.NewTicker(hb_interval)
		for {
			select {
			case <-ctx.Done():
				break
			case <-hb_ticker.C:
				msg := make(map[string]interface{})
				msg["op"] = 1
				if c.seq == 0 {
					msg["d"] = nil
				} else {
					msg["d"] = c.seq
				}
				c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
				log.Printf("Sending Heartbeat: %+v\n", msg)
				err = c.conn.WriteJSON(msg)
				if err != nil {
					c.conn.Close()
					return err
				}
			}
		}
	case 11:
		log.Println("Heartbeat ACK!")
	default:
		log.Printf("%d: %+v\n", op.Int(), data)
	}

	return nil
}

func (c *Connection) GetGateway() (string, error) {
	req, err := http.NewRequest("GET", GATEWAYBOTURL, nil)
	if err != nil {
		return "", err
	}

	req.Header.Add("Authorization", "Bot "+c.token)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	b, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return gjson.GetBytes(b, "url").String(), err
}