gavalink/node.go

package gavalink

import (
	"encoding/json"
	"fmt"
	"github.com/valyala/fastjson"
	"net/http"
	"strconv"

	"github.com/gorilla/websocket"
)

// NodeConfig configures a Lavalink Node
type NodeConfig struct {
	// REST is the host where Lavalink's REST server runs
	//
	// This value is expected without a trailing slash, e.g. like
	// `http://localhost:2333`
	REST string
	// WebSocket is the host where Lavalink's WebSocket server runs
	//
	// This value is expected without a trailing slash, e.g. like
	// `http://localhost:8012`
	WebSocket string
	// Password is the expected Authorization header for the Node
	Password string
}

// Node wraps a Lavalink Node
type Node struct {
	config  NodeConfig
	load    float32
	manager *Lavalink
	wsConn  *websocket.Conn
	client  *http.Client
}

func (node *Node) open() error {
	header := http.Header{}

	header.Set("Authorization", node.config.Password)
	header.Set("Num-Shards", node.manager.shards)
	header.Set("User-Id", node.manager.userID)

	ws, resp, err := websocket.DefaultDialer.Dial(node.config.WebSocket, header)

	if err != nil {
		return err
	}

	vstr := resp.Header.Get("Lavalink-Major-Version")

	v, err := strconv.Atoi(vstr)

	if err != nil {
		return err
	}

	if v < 3 {
		return errInvalidVersion
	}

	node.wsConn = ws
	go node.listen()

	Log.Println("node", node.config.WebSocket, "opened")

	return nil
}

func (node *Node) stop() {
	// someone already stopped this
	if node.wsConn == nil {
		return
	}

	_ = node.wsConn.Close()
}

func (node *Node) listen() {
	var p fastjson.Parser

	for {
		msgType, msg, err := node.wsConn.ReadMessage()

		if err != nil {
			Log.Println(err)
			// try to reconnect
			oerr := node.open()

			if oerr != nil {
				Log.Println("node", node.config.WebSocket, "failed and could not reconnect, destroying.", err, oerr)
				node.manager.removeNode(node)
				return
			}

			Log.Println("node", node.config.WebSocket, "reconnected")
			return
		}

		if msgType != websocket.TextMessage {
			continue
		}

		v, err := p.ParseBytes(msg)

		if err != nil {
			continue
		}

		node.onEvent(v)
	}
}

func (node *Node) onEvent(v *fastjson.Value) error {
	op := jsonStringValue(v, "op")

	switch op {
	case opPlayerUpdate:
		player, err := node.manager.GetPlayer(jsonStringValue(v, "guildId"))
		if err != nil {
			return err
		}

		player.time = v.GetInt("state", "time")
		player.position = v.GetInt("state", "position")
	case opEvent:
		player, err := node.manager.GetPlayer(jsonStringValue(v, "guildId"))
		if err != nil {
			return err
		}

		track := jsonStringValue(v, "track")

		switch jsonStringValue(v, "type") {
		case eventTrackEnd:
			player.track = ""
			err = player.handler.OnTrackEnd(player, track, jsonStringValue(v, "reason"))
		case eventTrackException:
			err = player.handler.OnTrackException(player, track, jsonStringValue(v, "reason"))
		case eventTrackStuck:
			err = player.handler.OnTrackStuck(player, track, v.GetInt("thresholdMs"))
		}

		return err
	case opVoiceProcessed:
		player, err := node.manager.GetPlayer(jsonStringValue(v, "guildId"))

		if err != nil {
			return err
		}

		track := jsonStringValue(v, "track")

		data := &VoiceProcessingData{
			URL:  fmt.Sprintf("%s/audio/%s", node.config.REST, track),
			File: track,
		}

		return player.handler.OnVoiceProcessed(player, data, v.GetBool("hotword"), v.GetBool("override"))
	case opStats:
		node.load = float32(v.GetFloat64("cpu", "lavalinkLoad"))
	default:
		return errUnknownPayload
	}

	return nil
}

// CreatePlayer creates an audio player on this node
func (node *Node) CreatePlayer(guildID string, sessionID string, event VoiceServerUpdate, handler EventHandler) (*Player, error) {
	msg := voiceUpdateMessage{
		Op:        opVoiceUpdate,
		GuildID:   guildID,
		SessionID: sessionID,
		Event:     &event,
	}

	err := node.writeMessage(msg)

	if err != nil {
		return nil, err
	}

	player := &Player{
		guildID:               guildID,
		manager:               node.manager,
		node:                  node,
		handler:               handler,
		vol:                   100,
		lastVoiceServerUpdate: event,
	}

	node.manager.players[guildID] = player

	return player, nil
}

// LoadTracks queries lavalink to return a Tracks object
//
// query should be a valid Lavaplayer query, including but not limited to:
// - A direct media URI
// - A direct Youtube /watch URI
// - A search query, prefixed with ytsearch: or scsearch:
//
// See the Lavaplayer Source Code for all valid options.
func (node *Node) LoadTracks(query string) (*Tracks, error) {
	url := fmt.Sprintf("%s/loadtracks?identifier=%s", node.config.REST, query)

	req, err := http.NewRequest(http.MethodGet, url, nil)

	if err != nil {
		return nil, err
	}

	req.Header.Set("Authorization", node.config.Password)

	resp, err := node.client.Do(req)

	if err != nil {
		return nil, err
	}

	tracks := new(Tracks)

	if err := json.NewDecoder(resp.Body).Decode(&tracks); err != nil {
		return nil, err
	}

	return tracks, nil
}

func (node *Node) writeMessage(v interface{}) error {
	return node.wsConn.WriteJSON(v)
}