From 3a9f33fc078240ba26dda76a83589bb8dbd605cf Mon Sep 17 00:00:00 2001
From: Christopher F <foxbot@protonmail.com>
Date: Wed, 22 Aug 2018 19:54:57 -0400
Subject: [PATCH] fix: minor fixes to make lavalink work

---
 example/main.go | 142 ++++++++++++++++++++++++++++++++++++++++++++++++
 lavalink.go     |  26 +++++----
 node.go         |   5 +-
 3 files changed, 159 insertions(+), 14 deletions(-)
 create mode 100644 example/main.go

diff --git a/example/main.go b/example/main.go
new file mode 100644
index 0000000..3271bbf
--- /dev/null
+++ b/example/main.go
@@ -0,0 +1,142 @@
+package main
+
+import (
+	"flag"
+	"log"
+	"os"
+	"os/signal"
+	"strconv"
+	"syscall"
+
+	"github.com/foxbot/gavalink"
+
+	"github.com/bwmarrin/discordgo"
+)
+
+var token string
+var lavalink *gavalink.Lavalink
+var player *gavalink.Player
+
+func init() {
+	flag.StringVar(&token, "token", "", "token=unprefixed token")
+}
+
+func main() {
+	flag.Parse()
+
+	if token == "" {
+		panic("no token specified!")
+	}
+	token = "Bot " + token
+
+	dg, err := discordgo.New(token)
+	if err != nil {
+		panic(err)
+	}
+
+	dg.AddHandler(ready)
+	dg.AddHandler(messageCreate)
+	dg.AddHandler(voiceServerUpdate)
+
+	err = dg.Open()
+	if err != nil {
+		panic(err)
+	}
+
+	sc := make(chan os.Signal, 1)
+	signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
+	<-sc
+
+	dg.Close()
+}
+
+func ready(s *discordgo.Session, event *discordgo.Ready) {
+	log.Println("discordgo ready!")
+	s.UpdateStatus(0, "gavalink")
+
+	lavalink = gavalink.NewLavalink("1", event.User.ID)
+
+	err := lavalink.AddNodes(gavalink.NodeConfig{
+		REST:      "http://localhost:2333",
+		WebSocket: "ws://localhost:2334",
+		Password:  "youshallnotpass",
+	})
+	if err != nil {
+		log.Println(err)
+	}
+}
+
+func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
+	if m.Author.ID == s.State.User.ID {
+		return
+	}
+
+	if m.Content == "~>>join" {
+		c, err := s.State.Channel(m.ChannelID)
+		if err != nil {
+			log.Println("fail find channel")
+			return
+		}
+
+		g, err := s.State.Guild(c.GuildID)
+		if err != nil {
+			log.Println("fail find guild")
+			return
+		}
+
+		for _, vs := range g.VoiceStates {
+			if vs.UserID == m.Author.ID {
+				log.Println("trying to connect to channel")
+				_, err := s.ChannelVoiceJoin(c.GuildID, vs.ChannelID, false, false)
+				if err != nil {
+					log.Println(err)
+				} else {
+					log.Println("channel voice join succeeded")
+				}
+			}
+		}
+	} else if m.Content == "~>>playtest" {
+		node, err := lavalink.BestNode()
+		if err != nil {
+			log.Println(err)
+			return
+		}
+		tracks, err := node.LoadTracks("https://www.youtube.com/watch?v=Wl2Q_MceIUc")
+		if err != nil {
+			log.Println(err)
+			return
+		}
+		if tracks.Type != gavalink.TrackLoaded {
+			log.Println("weird tracks type", tracks.Type)
+		}
+		track := tracks.Tracks[0].Data
+		err = player.Play(track)
+		if err != nil {
+			log.Println(err)
+		}
+	}
+}
+
+func voiceServerUpdate(s *discordgo.Session, event *discordgo.VoiceServerUpdate) {
+	log.Println("received VSU")
+	node, err := lavalink.BestNode()
+	if err != nil {
+		log.Println(err)
+		return
+	}
+
+	gid, _ := strconv.Atoi(event.GuildID)
+	vsu := gavalink.VoiceServerUpdate{
+		Endpoint: event.Endpoint,
+		GuildID:  gid,
+		Token:    event.Token,
+	}
+	handler := new(gavalink.DummyEventHandler)
+	player, err = node.CreatePlayer(event.GuildID, s.State.SessionID, vsu, handler)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+
+	panic("discordgo: cancel event")
+}
diff --git a/lavalink.go b/lavalink.go
index 5031e12..8c5f8ad 100644
--- a/lavalink.go
+++ b/lavalink.go
@@ -6,10 +6,8 @@ import (
 
 // Lavalink manages a connection to Lavalink Nodes
 type Lavalink struct {
-	// Shards is the total number of shards the bot is running
-	Shards int
-	// UserID is the Discord User ID of the bot
-	UserID int
+	shards string
+	userID string
 
 	nodes   []Node
 	players map[string]*Player
@@ -26,27 +24,31 @@ var (
 )
 
 // NewLavalink creates a new Lavalink manager
-func NewLavalink(shards int, userID int) *Lavalink {
+func NewLavalink(shards string, userID string) *Lavalink {
 	return &Lavalink{
-		Shards:  shards,
-		UserID:  userID,
-		nodes:   make([]Node, 1),
+		shards: shards,
+		userID: userID,
+		/*		nodes:   make([]Node, 1),*/
 		players: make(map[string]*Player),
 	}
 }
 
 // AddNodes adds a node to the Lavalink manager
-func (lavalink *Lavalink) AddNodes(nodeConfigs ...NodeConfig) {
+func (lavalink *Lavalink) AddNodes(nodeConfigs ...NodeConfig) error {
 	nodes := make([]Node, len(nodeConfigs))
 	for i, c := range nodeConfigs {
 		n := Node{
-			config: c,
-			shards: lavalink.Shards,
-			userID: lavalink.UserID,
+			config:  c,
+			manager: lavalink,
+		}
+		err := n.open()
+		if err != nil {
+			return err
 		}
 		nodes[i] = n
 	}
 	lavalink.nodes = append(lavalink.nodes, nodes...)
+	return nil
 }
 
 // RemoveNode removes a node from the manager
diff --git a/node.go b/node.go
index 65e3111..f360d6d 100644
--- a/node.go
+++ b/node.go
@@ -30,8 +30,6 @@ type NodeConfig struct {
 // Node wraps a Lavalink Node
 type Node struct {
 	config  NodeConfig
-	shards  int
-	userID  int
 	manager *Lavalink
 	wsConn  *websocket.Conn
 }
@@ -39,6 +37,8 @@ type Node struct {
 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 {
@@ -73,6 +73,7 @@ func (node *Node) listen() {
 	for {
 		msgType, msg, err := node.wsConn.ReadMessage()
 		if err != nil {
+			log.Println(err)
 			// try to reconnect
 			oerr := node.open()
 			if oerr != nil {