<commit_before>// Copyright 2014 docker-cluster authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mongodb

import (
	"time"

	"github.com/tsuru/docker-cluster/cluster"
	"github.com/tsuru/docker-cluster/storage"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type mongodbStorage struct {
	session *mgo.Session
	dbName  string
}

func (s *mongodbStorage) StoreContainer(container, host string) error {
	coll := s.getColl("containers")
	defer coll.Database.Session.Close()
	_, err := coll.UpsertId(container, bson.M{"$set": bson.M{"host": host}})
	return err
}

func (s *mongodbStorage) RetrieveContainer(container string) (string, error) {
	coll := s.getColl("containers")
	defer coll.Database.Session.Close()
	dbContainer := struct {
		Host string
	}{}
	err := coll.Find(bson.M{"_id": container}).One(&dbContainer)
	if err != nil {
		if err == mgo.ErrNotFound {
			return "", storage.ErrNoSuchContainer
		}
		return "", err
	}
	return dbContainer.Host, nil
}

func (s *mongodbStorage) RemoveContainer(container string) error {
	coll := s.getColl("containers")
	defer coll.Database.Session.Close()
	return coll.Remove(bson.M{"_id": container})
}

func (s *mongodbStorage) RetrieveContainers() ([]cluster.Container, error) {
	coll := s.getColl("containers")
	defer coll.Database.Session.Close()
	var containers []cluster.Container
	err := coll.Find(nil).All(&containers)
	return containers, err
}

func (s *mongodbStorage) StoreImage(repo, id, host string) error {
	coll := s.getColl("images_history")
	defer coll.Database.Session.Close()
	_, err := coll.UpsertId(repo, bson.M{
		"$addToSet": bson.M{"history": bson.D([]bson.DocElem{
			// Order is important for $addToSet!
			bson.DocElem{Name: "node", Value: host}, bson.DocElem{Name: "imageid", Value: id},
		})},
		"$set": bson.M{"lastnode": host, "lastid": id},
	})
	return err
}

func (s *mongodbStorage) RetrieveImage(repo string) (cluster.Image, error) {
	coll := s.getColl("images_history")
	defer coll.Database.Session.Close()
	var image cluster.Image
	err := coll.FindId(repo).One(&image)
	if err != nil {
		if err == mgo.ErrNotFound {
			return image, storage.ErrNoSuchImage
		}
		return image, err
	}
	if len(image.History) == 0 {
		return image, storage.ErrNoSuchImage
	}
	return image, nil
}

func (s *mongodbStorage) RemoveImage(repo, id, host string) error {
	coll := s.getColl("images_history")
	defer coll.Database.Session.Close()
	return coll.UpdateId(repo, bson.M{"$pull": bson.M{"history": bson.M{"node": host, "imageid": id}}})
}

func (s *mongodbStorage) RetrieveImages() ([]cluster.Image, error) {
	coll := s.getColl("images_history")
	defer coll.Database.Session.Close()
	var images []cluster.Image
	err := coll.Find(nil).All(&images)
	return images, err
}

func (s *mongodbStorage) StoreNode(node cluster.Node) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	err := coll.Insert(node)
	if mgo.IsDup(err) {
		return storage.ErrDuplicatedNodeAddress
	}
	return err
}

func (s *mongodbStorage) LockNodeForHealing(address string, isFailure bool, timeout time.Duration) (bool, error) {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	now := time.Now().UTC()
	until := now.Add(timeout)
	setOperation := bson.M{"$set": bson.M{"healing": bson.M{"lockeduntil": until, "isfailure": isFailure}}}
	err := coll.Update(
		bson.M{"_id": address, "healing.lockeduntil": nil},
		setOperation)
	if err == mgo.ErrNotFound {
		var dbNode cluster.Node
		err = coll.Find(bson.M{"_id": address}).One(&dbNode)
		if dbNode.Healing.LockedUntil.After(now) {
			return false, nil
		}
		err = coll.Update(bson.M{
			"_id": address,
			"healing.lockeduntil": dbNode.Healing.LockedUntil,
		}, setOperation)
		if err == mgo.ErrNotFound {
			return false, nil
		}
	}
	return err == nil, err
}

func (s *mongodbStorage) ExtendNodeLock(address string, timeout time.Duration) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	now := time.Now().UTC()
	until := now.Add(timeout)
	return coll.Update(
		bson.M{"_id": address},
		bson.M{"$set": bson.M{"healing.lockeduntil": until}})
}

func (s *mongodbStorage) UnlockNode(address string) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	return coll.Update(
		bson.M{"_id": address},
		bson.M{"$set": bson.M{"healing": nil}})
}

func (s *mongodbStorage) RetrieveNodesByMetadata(metadata map[string]string) ([]cluster.Node, error) {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	query := bson.M{}
	for key, value := range metadata {
		query["metadata."+key] = value
	}
	var nodes []cluster.Node
	err := coll.Find(query).All(&nodes)
	if err != nil {
		return nil, err
	}
	return nodes, nil
}

func (s *mongodbStorage) RetrieveNodes() ([]cluster.Node, error) {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	var nodes []cluster.Node
	err := coll.Find(nil).All(&nodes)
	if err != nil {
		return nil, err
	}
	return nodes, nil
}

func (s *mongodbStorage) RetrieveNode(address string) (cluster.Node, error) {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	var node cluster.Node
	err := coll.FindId(address).One(&node)
	if err == mgo.ErrNotFound {
		return cluster.Node{}, storage.ErrNoSuchNode
	}
	return node, err
}

func (s *mongodbStorage) UpdateNode(node cluster.Node) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	err := coll.UpdateId(node.Address, node)
	if err == mgo.ErrNotFound {
		return storage.ErrNoSuchNode
	}
	return err
}

func (s *mongodbStorage) RemoveNode(address string) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	err := coll.Remove(bson.M{"_id": address})
	if err == mgo.ErrNotFound {
		return storage.ErrNoSuchNode
	}
	return err
}

func (s *mongodbStorage) getColl(name string) *mgo.Collection {
	session := s.session.Clone()
	return session.DB(s.dbName).C(name)
}

func Mongodb(addr, dbName string) (cluster.Storage, error) {
	session, err := mgo.Dial(addr)
	if err != nil {
		return nil, err
	}
	storage := mongodbStorage{
		session: session,
		dbName:  dbName,
	}
	return &storage, nil
}
<commit_msg>storage/mongodb: fail fast and more aggressive timeout for connections<commit_after>// Copyright 2014 docker-cluster authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mongodb

import (
	"time"

	"github.com/tsuru/docker-cluster/cluster"
	"github.com/tsuru/docker-cluster/storage"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type mongodbStorage struct {
	session *mgo.Session
	dbName  string
}

func (s *mongodbStorage) StoreContainer(container, host string) error {
	coll := s.getColl("containers")
	defer coll.Database.Session.Close()
	_, err := coll.UpsertId(container, bson.M{"$set": bson.M{"host": host}})
	return err
}

func (s *mongodbStorage) RetrieveContainer(container string) (string, error) {
	coll := s.getColl("containers")
	defer coll.Database.Session.Close()
	dbContainer := struct {
		Host string
	}{}
	err := coll.Find(bson.M{"_id": container}).One(&dbContainer)
	if err != nil {
		if err == mgo.ErrNotFound {
			return "", storage.ErrNoSuchContainer
		}
		return "", err
	}
	return dbContainer.Host, nil
}

func (s *mongodbStorage) RemoveContainer(container string) error {
	coll := s.getColl("containers")
	defer coll.Database.Session.Close()
	return coll.Remove(bson.M{"_id": container})
}

func (s *mongodbStorage) RetrieveContainers() ([]cluster.Container, error) {
	coll := s.getColl("containers")
	defer coll.Database.Session.Close()
	var containers []cluster.Container
	err := coll.Find(nil).All(&containers)
	return containers, err
}

func (s *mongodbStorage) StoreImage(repo, id, host string) error {
	coll := s.getColl("images_history")
	defer coll.Database.Session.Close()
	_, err := coll.UpsertId(repo, bson.M{
		"$addToSet": bson.M{"history": bson.D([]bson.DocElem{
			// Order is important for $addToSet!
			bson.DocElem{Name: "node", Value: host}, bson.DocElem{Name: "imageid", Value: id},
		})},
		"$set": bson.M{"lastnode": host, "lastid": id},
	})
	return err
}

func (s *mongodbStorage) RetrieveImage(repo string) (cluster.Image, error) {
	coll := s.getColl("images_history")
	defer coll.Database.Session.Close()
	var image cluster.Image
	err := coll.FindId(repo).One(&image)
	if err != nil {
		if err == mgo.ErrNotFound {
			return image, storage.ErrNoSuchImage
		}
		return image, err
	}
	if len(image.History) == 0 {
		return image, storage.ErrNoSuchImage
	}
	return image, nil
}

func (s *mongodbStorage) RemoveImage(repo, id, host string) error {
	coll := s.getColl("images_history")
	defer coll.Database.Session.Close()
	return coll.UpdateId(repo, bson.M{"$pull": bson.M{"history": bson.M{"node": host, "imageid": id}}})
}

func (s *mongodbStorage) RetrieveImages() ([]cluster.Image, error) {
	coll := s.getColl("images_history")
	defer coll.Database.Session.Close()
	var images []cluster.Image
	err := coll.Find(nil).All(&images)
	return images, err
}

func (s *mongodbStorage) StoreNode(node cluster.Node) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	err := coll.Insert(node)
	if mgo.IsDup(err) {
		return storage.ErrDuplicatedNodeAddress
	}
	return err
}

func (s *mongodbStorage) LockNodeForHealing(address string, isFailure bool, timeout time.Duration) (bool, error) {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	now := time.Now().UTC()
	until := now.Add(timeout)
	setOperation := bson.M{"$set": bson.M{"healing": bson.M{"lockeduntil": until, "isfailure": isFailure}}}
	err := coll.Update(
		bson.M{"_id": address, "healing.lockeduntil": nil},
		setOperation)
	if err == mgo.ErrNotFound {
		var dbNode cluster.Node
		err = coll.Find(bson.M{"_id": address}).One(&dbNode)
		if dbNode.Healing.LockedUntil.After(now) {
			return false, nil
		}
		err = coll.Update(bson.M{
			"_id": address,
			"healing.lockeduntil": dbNode.Healing.LockedUntil,
		}, setOperation)
		if err == mgo.ErrNotFound {
			return false, nil
		}
	}
	return err == nil, err
}

func (s *mongodbStorage) ExtendNodeLock(address string, timeout time.Duration) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	now := time.Now().UTC()
	until := now.Add(timeout)
	return coll.Update(
		bson.M{"_id": address},
		bson.M{"$set": bson.M{"healing.lockeduntil": until}})
}

func (s *mongodbStorage) UnlockNode(address string) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	return coll.Update(
		bson.M{"_id": address},
		bson.M{"$set": bson.M{"healing": nil}})
}

func (s *mongodbStorage) RetrieveNodesByMetadata(metadata map[string]string) ([]cluster.Node, error) {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	query := bson.M{}
	for key, value := range metadata {
		query["metadata."+key] = value
	}
	var nodes []cluster.Node
	err := coll.Find(query).All(&nodes)
	if err != nil {
		return nil, err
	}
	return nodes, nil
}

func (s *mongodbStorage) RetrieveNodes() ([]cluster.Node, error) {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	var nodes []cluster.Node
	err := coll.Find(nil).All(&nodes)
	if err != nil {
		return nil, err
	}
	return nodes, nil
}

func (s *mongodbStorage) RetrieveNode(address string) (cluster.Node, error) {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	var node cluster.Node
	err := coll.FindId(address).One(&node)
	if err == mgo.ErrNotFound {
		return cluster.Node{}, storage.ErrNoSuchNode
	}
	return node, err
}

func (s *mongodbStorage) UpdateNode(node cluster.Node) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	err := coll.UpdateId(node.Address, node)
	if err == mgo.ErrNotFound {
		return storage.ErrNoSuchNode
	}
	return err
}

func (s *mongodbStorage) RemoveNode(address string) error {
	coll := s.getColl("nodes")
	defer coll.Database.Session.Close()
	err := coll.Remove(bson.M{"_id": address})
	if err == mgo.ErrNotFound {
		return storage.ErrNoSuchNode
	}
	return err
}

func (s *mongodbStorage) getColl(name string) *mgo.Collection {
	session := s.session.Copy()
	return session.DB(s.dbName).C(name)
}

func Mongodb(addr, dbName string) (cluster.Storage, error) {
	dialInfo, err := mgo.ParseURL(addr)
	if err != nil {
		return nil, err
	}
	dialInfo.FailFast = true
	session, err := mgo.DialWithInfo(dialInfo)
	if err != nil {
		return nil, err
	}
	session.SetSyncTimeout(10 * time.Second)
	session.SetSocketTimeout(1 * time.Minute)
	session.SetMode(mgo.Monotonic, true)
	storage := mongodbStorage{
		session: session,
		dbName:  dbName,
	}
	return &storage, nil
}
<|endoftext|>


<commit_before>package web

import (
	"database/sql"
	"encoding/json"
	"log"
	"time"

	"github.com/barracudanetworks/wbd/database"
	"github.com/gorilla/websocket"
)

const (
	writeWait = 10 * time.Second

	pongWait = 60 * time.Second
	pingWait = (pongWait * 9) / 10

	urlPollWait = pongWait / 2

	maxMessageSize = 512
)

type websocketMessage struct {
	Action string      `json:"action"`
	Data   interface{} `json:"data"`
}

// The map is a bit weird in that it's pointer->string, but it's kind of a cheap
// hack around issues with overwriting a connection that has the same client name
type websocketHub struct {
	broadcast   chan *websocketMessage
	register    chan *websocketClient
	unregister  chan *websocketClient
	connections map[*websocketClient]string
}

var hub = websocketHub{
	broadcast:   make(chan *websocketMessage),
	register:    make(chan *websocketClient),
	unregister:  make(chan *websocketClient),
	connections: make(map[*websocketClient]string),
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func (h *websocketHub) run(a *App) {
	db := a.Database

	ticker := time.NewTicker(urlPollWait)
	defer func() {
		ticker.Stop()
	}()

	for {
		select {
		// Save connection to hub
		case c := <-h.register:
			log.Printf("Added client '%s' to the hub", c.Id)

			if !c.Generic {
				// get client info and touch last ping
				client, err := db.GetClient(c.Id)
				switch {
				case err == sql.ErrNoRows:
					log.Printf("Unknown client, creating record")
					if err := db.InsertClient(c.Id, c.IpAddress); err != nil {
						log.Fatal(err)
					}
				case err != nil:
					log.Fatal(err)
				default:
					log.Printf("Client last seen %s from %s", client.LastPing, client.IpAddress)
					db.TouchClient(client.Identifier)
				}
			} else {
				log.Printf("Not attempting to track generic client")
			}

			h.connections[c] = c.Id

		// Remove connection from hub
		case c := <-h.unregister:
			if _, ok := h.connections[c]; ok {
				log.Printf("Removing client '%s'", c.Id)
				h.CloseConnection(c)
			}

		// Broadcast messages to clients
		case m := <-h.broadcast:
			for c := range h.connections {
				select {
				case c.send <- m:
					log.Printf("Sent message to client '%s', type '%s'", c.Id, m.Action)
				default:
					h.CloseConnection(c)
				}
			}

		// Send out URL updates every so often
		case <-ticker.C:
			log.Print("Polling for URL changes in database")

			// Send the JSON to all connected clients
			for c := range h.connections {
				urls, err := db.FetchUrlsByClientId(c.Id)
				if err != nil {
					log.Fatal(err)
				}

				urlWm, err := urlUpdateMessage(urls)
				if err != nil {
					log.Fatal(err)
				}

				select {
				case c.send <- urlWm:
					log.Printf("Sent updated URL list to client '%s'", c.Id)
				default:
					h.CloseConnection(c)
				}
			}
		}
	}
}

func (h *websocketHub) CloseConnection(c *websocketClient) {
	log.Printf("Closing connection to client '%s'", c.Id)
	close(c.send)
	delete(h.connections, c)
}

type websocketClient struct {
	Id         string
	IpAddress  string
	Controller bool
	Generic    bool

	ws   *websocket.Conn
	send chan *websocketMessage
}

func (c *websocketClient) write(mt int, payload []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.ws.WriteMessage(mt, payload)
}

func (c *websocketClient) writePump() {
	ticker := time.NewTicker(pingWait)
	defer func() {
		ticker.Stop()
		c.ws.Close()
	}()

	for {
		select {
		// If there's a message in the send queue, write to the socket
		case message, ok := <-c.send:
			if !ok {
				c.write(websocket.CloseMessage, []byte{})
				return
			}

			json, err := json.Marshal(message)
			if err != nil {
				log.Print(err)
				return
			}

			if err := c.write(websocket.TextMessage, json); err != nil {
				log.Print(err)
				return
			}

		// Send ping on timer
		case <-ticker.C:
			if err := c.write(websocket.PingMessage, []byte{}); err != nil {
				log.Print(err)
				return
			}
		}
	}
}

func (c *websocketClient) readPump(db *database.Database) {
	defer func() {
		hub.unregister <- c
		c.ws.Close()
	}()

	c.ws.SetReadLimit(maxMessageSize)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })

	for {
		_, message, err := c.ws.ReadMessage()
		if err != nil {
			log.Printf("Connection closed to client '%s'", c.Id)
			break
		}

		var wm websocketMessage
		if err := json.Unmarshal(message, &wm); err != nil {
			log.Print(err)
			break
		}

		// Respond to a requested action by the client
		switch wm.Action {
		case "flagController":
			log.Printf("Client '%s' flagged as a controller", c.Id)
			c.Controller = true

			urls, err := db.FetchUrlsByClientId(c.Id)
			if err != nil {
				log.Fatal(err)
			}

			urlWm, err := urlUpdateMessage(urls)
			if err != nil {
				log.Fatal(err)
			}

			clientWm := hub.clientUpdateMessage()

			c.send <- urlWm
			c.send <- clientWm
		case "sendUrls":
			log.Printf("Client '%s' requested URLs", c.Id)

			urls, err := db.FetchUrlsByClientId(c.Id)
			if err != nil {
				log.Fatal(err)
			}

			urlWm, err := urlUpdateMessage(urls)
			if err != nil {
				log.Fatal(err)
			}

			c.send <- urlWm
		case "sendClients":
			log.Printf("Client '%s' requested clients", c.Id)

			clientWm := hub.clientUpdateMessage()

			c.send <- clientWm
		default:
			log.Printf("Unknown action %s from client '%s'", wm.Action, c.Id)
		}
	}
}

func urlUpdateMessage(urls []string) (wm *websocketMessage, err error) {
	wm = &websocketMessage{
		Action: "updateUrls",
		Data: struct {
			URLs []string `json:"urls"`
		}{
			urls,
		},
	}

	return
}

func (h *websocketHub) GetClients() (clients []string) {
	for c := range h.connections {
		if c.Id != "" {
			clients = append(clients, c.Id)
		}
	}

	return
}

func (h *websocketHub) clientUpdateMessage() (wm *websocketMessage) {
	clients := h.GetClients()

	wm = &websocketMessage{
		Action: "updateClients",
		Data: struct {
			Clients []string `json:"clients"`
		}{
			clients,
		},
	}

	return
}
<commit_msg>Update client IP address on websocket connection<commit_after>package web

import (
	"database/sql"
	"encoding/json"
	"log"
	"time"

	"github.com/barracudanetworks/wbd/database"
	"github.com/gorilla/websocket"
)

const (
	writeWait = 10 * time.Second

	pongWait = 60 * time.Second
	pingWait = (pongWait * 9) / 10

	urlPollWait = pongWait / 2

	maxMessageSize = 512
)

type websocketMessage struct {
	Action string      `json:"action"`
	Data   interface{} `json:"data"`
}

// The map is a bit weird in that it's pointer->string, but it's kind of a cheap
// hack around issues with overwriting a connection that has the same client name
type websocketHub struct {
	broadcast   chan *websocketMessage
	register    chan *websocketClient
	unregister  chan *websocketClient
	connections map[*websocketClient]string
}

var hub = websocketHub{
	broadcast:   make(chan *websocketMessage),
	register:    make(chan *websocketClient),
	unregister:  make(chan *websocketClient),
	connections: make(map[*websocketClient]string),
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func (h *websocketHub) run(a *App) {
	db := a.Database

	ticker := time.NewTicker(urlPollWait)
	defer func() {
		ticker.Stop()
	}()

	for {
		select {
		// Save connection to hub
		case c := <-h.register:
			log.Printf("Added client '%s' to the hub", c.Id)

			if !c.Generic {
				// get client info and touch last ping
				client, err := db.GetClient(c.Id)
				switch {
				case err == sql.ErrNoRows:
					log.Printf("Unknown client, creating record")
					if err := db.InsertClient(c.Id, c.IpAddress); err != nil {
						log.Fatal(err)
					}
				case err != nil:
					log.Fatal(err)
				default:
					log.Printf("Client last seen %s from %s", client.LastPing, client.IpAddress)

					db.TouchClient(client.Identifier)
					db.SetClientIpAddress(client.Identifier, c.IpAddress)
				}
			} else {
				log.Printf("Not attempting to track generic client")
			}

			h.connections[c] = c.Id

		// Remove connection from hub
		case c := <-h.unregister:
			if _, ok := h.connections[c]; ok {
				log.Printf("Removing client '%s'", c.Id)
				h.CloseConnection(c)
			}

		// Broadcast messages to clients
		case m := <-h.broadcast:
			for c := range h.connections {
				select {
				case c.send <- m:
					log.Printf("Sent message to client '%s', type '%s'", c.Id, m.Action)
				default:
					h.CloseConnection(c)
				}
			}

		// Send out URL updates every so often
		case <-ticker.C:
			log.Print("Polling for URL changes in database")

			// Send the JSON to all connected clients
			for c := range h.connections {
				urls, err := db.FetchUrlsByClientId(c.Id)
				if err != nil {
					log.Fatal(err)
				}

				urlWm, err := urlUpdateMessage(urls)
				if err != nil {
					log.Fatal(err)
				}

				select {
				case c.send <- urlWm:
					log.Printf("Sent updated URL list to client '%s'", c.Id)
				default:
					h.CloseConnection(c)
				}
			}
		}
	}
}

func (h *websocketHub) CloseConnection(c *websocketClient) {
	log.Printf("Closing connection to client '%s'", c.Id)
	close(c.send)
	delete(h.connections, c)
}

type websocketClient struct {
	Id         string
	IpAddress  string
	Controller bool
	Generic    bool

	ws   *websocket.Conn
	send chan *websocketMessage
}

func (c *websocketClient) write(mt int, payload []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.ws.WriteMessage(mt, payload)
}

func (c *websocketClient) writePump() {
	ticker := time.NewTicker(pingWait)
	defer func() {
		ticker.Stop()
		c.ws.Close()
	}()

	for {
		select {
		// If there's a message in the send queue, write to the socket
		case message, ok := <-c.send:
			if !ok {
				c.write(websocket.CloseMessage, []byte{})
				return
			}

			json, err := json.Marshal(message)
			if err != nil {
				log.Print(err)
				return
			}

			if err := c.write(websocket.TextMessage, json); err != nil {
				log.Print(err)
				return
			}

		// Send ping on timer
		case <-ticker.C:
			if err := c.write(websocket.PingMessage, []byte{}); err != nil {
				log.Print(err)
				return
			}
		}
	}
}

func (c *websocketClient) readPump(db *database.Database) {
	defer func() {
		hub.unregister <- c
		c.ws.Close()
	}()

	c.ws.SetReadLimit(maxMessageSize)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })

	for {
		_, message, err := c.ws.ReadMessage()
		if err != nil {
			log.Printf("Connection closed to client '%s'", c.Id)
			break
		}

		var wm websocketMessage
		if err := json.Unmarshal(message, &wm); err != nil {
			log.Print(err)
			break
		}

		// Respond to a requested action by the client
		switch wm.Action {
		case "flagController":
			log.Printf("Client '%s' flagged as a controller", c.Id)
			c.Controller = true

			urls, err := db.FetchUrlsByClientId(c.Id)
			if err != nil {
				log.Fatal(err)
			}

			urlWm, err := urlUpdateMessage(urls)
			if err != nil {
				log.Fatal(err)
			}

			clientWm := hub.clientUpdateMessage()

			c.send <- urlWm
			c.send <- clientWm
		case "sendUrls":
			log.Printf("Client '%s' requested URLs", c.Id)

			urls, err := db.FetchUrlsByClientId(c.Id)
			if err != nil {
				log.Fatal(err)
			}

			urlWm, err := urlUpdateMessage(urls)
			if err != nil {
				log.Fatal(err)
			}

			c.send <- urlWm
		case "sendClients":
			log.Printf("Client '%s' requested clients", c.Id)

			clientWm := hub.clientUpdateMessage()

			c.send <- clientWm
		default:
			log.Printf("Unknown action %s from client '%s'", wm.Action, c.Id)
		}
	}
}

func urlUpdateMessage(urls []string) (wm *websocketMessage, err error) {
	wm = &websocketMessage{
		Action: "updateUrls",
		Data: struct {
			URLs []string `json:"urls"`
		}{
			urls,
		},
	}

	return
}

func (h *websocketHub) GetClients() (clients []string) {
	for c := range h.connections {
		if c.Id != "" {
			clients = append(clients, c.Id)
		}
	}

	return
}

func (h *websocketHub) clientUpdateMessage() (wm *websocketMessage) {
	clients := h.GetClients()

	wm = &websocketMessage{
		Action: "updateClients",
		Data: struct {
			Clients []string `json:"clients"`
		}{
			clients,
		},
	}

	return
}
<|endoftext|>


<commit_before>// reflectwalk is a package that allows you to "walk" complex structures
// similar to how you may "walk" a filesystem: visiting every element one
// by one and calling callback functions allowing you to handle and manipulate
// those elements.
package reflectwalk

import "reflect"

// PrimitiveWalker implementations are able to handle primitive values
// within complex structures. Primitive values are numbers, strings,
// booleans, funcs, chans.
//
// These primitive values are often members of more complex
// structures (slices, maps, etc.) that are walkable by other interfaces.
type PrimitiveWalker interface {
	Primitive(reflect.Value) error
}

// MapWalker implementations are able to handle individual elements
// found within a map structure.
type MapWalker interface {
	Map(m reflect.Value) error
	MapElem(m, k, v reflect.Value) error
}

// SliceWalker implementations are able to handle slice elements found
// within complex structures.
type SliceWalker interface {
	Slice(reflect.Value) error
	SliceElem(int, reflect.Value) error
}

// StructWalker is an interface that has methods that are called for
// structs when a Walk is done.
type StructWalker interface {
	Struct(reflect.Value) error
	StructField(reflect.StructField, reflect.Value) error
}

// EnterExitWalker implementations are notified before and after
// they walk deeper into complex structures (into struct fields,
// into slice elements, etc.)
type EnterExitWalker interface {
	Enter(Location) error
	Exit(Location) error
}

// PointerWalker implementations are notified when the value they're
// walking is a pointer or not. Pointer is called for _every_ value whether
// it is a pointer or not.
type PointerWalker interface {
	PointerEnter(bool) error
	PointerExit(bool) error
}

// Walk takes an arbitrary value and an interface and traverses the
// value, calling callbacks on the interface if they are supported.
// The interface should implement one or more of the walker interfaces
// in this package, such as PrimitiveWalker, StructWalker, etc.
func Walk(data, walker interface{}) (err error) {
	v := reflect.ValueOf(data)
	ew, ok := walker.(EnterExitWalker)
	if ok {
		err = ew.Enter(WalkLoc)
	}

	if err == nil {
		err = walk(v, walker)
	}

	if ok && err == nil {
		err = ew.Exit(WalkLoc)
	}

	return
}

func walk(v reflect.Value, w interface{}) (err error) {
	// We preserve the original value here because if it is an interface
	// type, we want to pass that directly into the walkPrimitive, so that
	// we can set it.
	originalV := v
	if v.Kind() == reflect.Interface {
		v = v.Elem()
	}

	// Determine if we're receiving a pointer and if so notify the walker.
	pointer := false
	if v.Kind() == reflect.Ptr {
		pointer = true
		v = reflect.Indirect(v)
	}
	if pw, ok := w.(PointerWalker); ok {
		if err = pw.PointerEnter(pointer); err != nil {
			return
		}

		defer func() {
			if err != nil {
				return
			}

			err = pw.PointerExit(pointer)
		}()
	}

	k := v.Kind()
	if k >= reflect.Int && k <= reflect.Complex128 {
		k = reflect.Int
	}

	switch k {
	// Primitives
	case reflect.Bool, reflect.Chan, reflect.Func, reflect.Int, reflect.String, reflect.Invalid:
		err = walkPrimitive(originalV, w)
		return
	case reflect.Map:
		err = walkMap(v, w)
		return
	case reflect.Slice:
		err = walkSlice(v, w)
		return
	case reflect.Struct:
		err = walkStruct(v, w)
		return
	default:
		panic("unsupported type: " + k.String())
	}
}

func walkMap(v reflect.Value, w interface{}) error {
	ew, ewok := w.(EnterExitWalker)
	if ewok {
		ew.Enter(Map)
	}

	if mw, ok := w.(MapWalker); ok {
		if err := mw.Map(v); err != nil {
			return err
		}
	}

	for _, k := range v.MapKeys() {
		kv := v.MapIndex(k)

		if mw, ok := w.(MapWalker); ok {
			if err := mw.MapElem(v, k, kv); err != nil {
				return err
			}
		}

		ew, ok := w.(EnterExitWalker)
		if ok {
			ew.Enter(MapKey)
		}

		if err := walk(k, w); err != nil {
			return err
		}

		if ok {
			ew.Exit(MapKey)
			ew.Enter(MapValue)
		}

		if err := walk(kv, w); err != nil {
			return err
		}

		if ok {
			ew.Exit(MapValue)
		}
	}

	if ewok {
		ew.Exit(Map)
	}

	return nil
}

func walkPrimitive(v reflect.Value, w interface{}) error {
	if pw, ok := w.(PrimitiveWalker); ok {
		return pw.Primitive(v)
	}

	return nil
}

func walkSlice(v reflect.Value, w interface{}) (err error) {
	ew, ok := w.(EnterExitWalker)
	if ok {
		ew.Enter(Slice)
	}

	if sw, ok := w.(SliceWalker); ok {
		if err := sw.Slice(v); err != nil {
			return err
		}
	}

	for i := 0; i < v.Len(); i++ {
		elem := v.Index(i)

		if sw, ok := w.(SliceWalker); ok {
			if err := sw.SliceElem(i, elem); err != nil {
				return err
			}
		}

		ew, ok := w.(EnterExitWalker)
		if ok {
			ew.Enter(SliceElem)
		}

		if err := walk(elem, w); err != nil {
			return err
		}

		if ok {
			ew.Exit(SliceElem)
		}
	}

	ew, ok = w.(EnterExitWalker)
	if ok {
		ew.Exit(Slice)
	}

	return nil
}

func walkStruct(v reflect.Value, w interface{}) (err error) {
	ew, ewok := w.(EnterExitWalker)
	if ewok {
		ew.Enter(Struct)
	}

	if sw, ok := w.(StructWalker); ok {
		if err = sw.Struct(v); err != nil {
			return
		}
	}

	vt := v.Type()
	for i := 0; i < vt.NumField(); i++ {
		sf := vt.Field(i)
		f := v.FieldByIndex([]int{i})

		if sw, ok := w.(StructWalker); ok {
			err = sw.StructField(sf, f)
			if err != nil {
				return
			}
		}

		ew, ok := w.(EnterExitWalker)
		if ok {
			ew.Enter(StructField)
		}

		err = walk(f, w)
		if err != nil {
			return
		}

		if ok {
			ew.Exit(StructField)
		}
	}

	if ewok {
		ew.Exit(Struct)
	}

	return nil
}
<commit_msg>Fix failing tests in copystructure, retain original behavior<commit_after>// reflectwalk is a package that allows you to "walk" complex structures
// similar to how you may "walk" a filesystem: visiting every element one
// by one and calling callback functions allowing you to handle and manipulate
// those elements.
package reflectwalk

import "reflect"

// PrimitiveWalker implementations are able to handle primitive values
// within complex structures. Primitive values are numbers, strings,
// booleans, funcs, chans.
//
// These primitive values are often members of more complex
// structures (slices, maps, etc.) that are walkable by other interfaces.
type PrimitiveWalker interface {
	Primitive(reflect.Value) error
}

// MapWalker implementations are able to handle individual elements
// found within a map structure.
type MapWalker interface {
	Map(m reflect.Value) error
	MapElem(m, k, v reflect.Value) error
}

// SliceWalker implementations are able to handle slice elements found
// within complex structures.
type SliceWalker interface {
	Slice(reflect.Value) error
	SliceElem(int, reflect.Value) error
}

// StructWalker is an interface that has methods that are called for
// structs when a Walk is done.
type StructWalker interface {
	Struct(reflect.Value) error
	StructField(reflect.StructField, reflect.Value) error
}

// EnterExitWalker implementations are notified before and after
// they walk deeper into complex structures (into struct fields,
// into slice elements, etc.)
type EnterExitWalker interface {
	Enter(Location) error
	Exit(Location) error
}

// PointerWalker implementations are notified when the value they're
// walking is a pointer or not. Pointer is called for _every_ value whether
// it is a pointer or not.
type PointerWalker interface {
	PointerEnter(bool) error
	PointerExit(bool) error
}

// Walk takes an arbitrary value and an interface and traverses the
// value, calling callbacks on the interface if they are supported.
// The interface should implement one or more of the walker interfaces
// in this package, such as PrimitiveWalker, StructWalker, etc.
func Walk(data, walker interface{}) (err error) {
	v := reflect.ValueOf(data)
	ew, ok := walker.(EnterExitWalker)
	if ok {
		err = ew.Enter(WalkLoc)
	}

	if err == nil {
		err = walk(v, walker)
	}

	if ok && err == nil {
		err = ew.Exit(WalkLoc)
	}

	return
}

func walk(v reflect.Value, w interface{}) (err error) {
	// Determine if we're receiving a pointer and if so notify the walker.
	// The logic here is convoluted but very important (tests will fail if
	// almost any part is changed). I will try to explain here.
	//
	// First, we check if the value is an interface, if so, we really need
	// to check the interface's VALUE to see whether it is a pointer (pointers
	// to interfaces are not allowed).
	//
	// Check whether the value is then an interface. If so, then set pointer
	// to true to notify the user.
	//
	// At this time, we also set "v" to be the dereferenced value. This is
	// because once we've unwrapped the pointer we want to use that value.
	pointer := false
	pointerV := v
	if pointerV.Kind() == reflect.Interface {
		pointerV = pointerV.Elem()
	}
	if pointerV.Kind() == reflect.Ptr {
		pointer = true
		v = reflect.Indirect(pointerV)
	}
	if pw, ok := w.(PointerWalker); ok {
		if err = pw.PointerEnter(pointer); err != nil {
			return
		}

		defer func() {
			if err != nil {
				return
			}

			err = pw.PointerExit(pointer)
		}()
	}

	// We preserve the original value here because if it is an interface
	// type, we want to pass that directly into the walkPrimitive, so that
	// we can set it.
	originalV := v
	if v.Kind() == reflect.Interface {
		v = v.Elem()
	}

	k := v.Kind()
	if k >= reflect.Int && k <= reflect.Complex128 {
		k = reflect.Int
	}

	switch k {
	// Primitives
	case reflect.Bool, reflect.Chan, reflect.Func, reflect.Int, reflect.String, reflect.Invalid:
		err = walkPrimitive(originalV, w)
		return
	case reflect.Map:
		err = walkMap(v, w)
		return
	case reflect.Slice:
		err = walkSlice(v, w)
		return
	case reflect.Struct:
		err = walkStruct(v, w)
		return
	default:
		panic("unsupported type: " + k.String())
	}
}

func walkMap(v reflect.Value, w interface{}) error {
	ew, ewok := w.(EnterExitWalker)
	if ewok {
		ew.Enter(Map)
	}

	if mw, ok := w.(MapWalker); ok {
		if err := mw.Map(v); err != nil {
			return err
		}
	}

	for _, k := range v.MapKeys() {
		kv := v.MapIndex(k)

		if mw, ok := w.(MapWalker); ok {
			if err := mw.MapElem(v, k, kv); err != nil {
				return err
			}
		}

		ew, ok := w.(EnterExitWalker)
		if ok {
			ew.Enter(MapKey)
		}

		if err := walk(k, w); err != nil {
			return err
		}

		if ok {
			ew.Exit(MapKey)
			ew.Enter(MapValue)
		}

		if err := walk(kv, w); err != nil {
			return err
		}

		if ok {
			ew.Exit(MapValue)
		}
	}

	if ewok {
		ew.Exit(Map)
	}

	return nil
}

func walkPrimitive(v reflect.Value, w interface{}) error {
	if pw, ok := w.(PrimitiveWalker); ok {
		return pw.Primitive(v)
	}

	return nil
}

func walkSlice(v reflect.Value, w interface{}) (err error) {
	ew, ok := w.(EnterExitWalker)
	if ok {
		ew.Enter(Slice)
	}

	if sw, ok := w.(SliceWalker); ok {
		if err := sw.Slice(v); err != nil {
			return err
		}
	}

	for i := 0; i < v.Len(); i++ {
		elem := v.Index(i)

		if sw, ok := w.(SliceWalker); ok {
			if err := sw.SliceElem(i, elem); err != nil {
				return err
			}
		}

		ew, ok := w.(EnterExitWalker)
		if ok {
			ew.Enter(SliceElem)
		}

		if err := walk(elem, w); err != nil {
			return err
		}

		if ok {
			ew.Exit(SliceElem)
		}
	}

	ew, ok = w.(EnterExitWalker)
	if ok {
		ew.Exit(Slice)
	}

	return nil
}

func walkStruct(v reflect.Value, w interface{}) (err error) {
	ew, ewok := w.(EnterExitWalker)
	if ewok {
		ew.Enter(Struct)
	}

	if sw, ok := w.(StructWalker); ok {
		if err = sw.Struct(v); err != nil {
			return
		}
	}

	vt := v.Type()
	for i := 0; i < vt.NumField(); i++ {
		sf := vt.Field(i)
		f := v.FieldByIndex([]int{i})

		if sw, ok := w.(StructWalker); ok {
			err = sw.StructField(sf, f)
			if err != nil {
				return
			}
		}

		ew, ok := w.(EnterExitWalker)
		if ok {
			ew.Enter(StructField)
		}

		err = walk(f, w)
		if err != nil {
			return
		}

		if ok {
			ew.Exit(StructField)
		}
	}

	if ewok {
		ew.Exit(Struct)
	}

	return nil
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package k6

import (
	"context"
	"fmt"
	"runtime"
	"testing"
	"time"

	"github.com/dop251/goja"
	"github.com/loadimpact/k6/js/common"
	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/metrics"
	"github.com/stretchr/testify/assert"
)

func TestFail(t *testing.T) {
	rt := goja.New()
	rt.Set("k6", common.Bind(rt, &K6{}, nil))
	_, err := common.RunString(rt, `k6.fail("blah")`)
	assert.EqualError(t, err, "GoError: blah")
}

func TestSleep(t *testing.T) {
	rt := goja.New()
	ctx, cancel := context.WithCancel(context.Background())
	rt.Set("k6", common.Bind(rt, &K6{}, &ctx))

	testdata := map[string]time.Duration{
		"1":   1 * time.Second,
		"1.0": 1 * time.Second,
		"0.5": 500 * time.Millisecond,
	}
	for name, d := range testdata {
		t.Run(name, func(t *testing.T) {
			startTime := time.Now()
			_, err := common.RunString(rt, `k6.sleep(1)`)
			endTime := time.Now()
			assert.NoError(t, err)
			assert.True(t, endTime.Sub(startTime) > d, "did not sleep long enough")
		})
	}

	t.Run("Cancel", func(t *testing.T) {
		dch := make(chan time.Duration)
		go func() {
			startTime := time.Now()
			_, err := common.RunString(rt, `k6.sleep(10)`)
			endTime := time.Now()
			assert.NoError(t, err)
			dch <- endTime.Sub(startTime)
		}()
		runtime.Gosched()
		time.Sleep(1 * time.Second)
		runtime.Gosched()
		cancel()
		runtime.Gosched()
		d := <-dch
		assert.True(t, d > 500*time.Millisecond, "did not sleep long enough")
		assert.True(t, d < 2*time.Second, "slept for too long!!")
	})
}

func TestGroup(t *testing.T) {
	root, err := lib.NewGroup("", nil)
	assert.NoError(t, err)

	rt := goja.New()
	state := &common.State{Group: root}

	ctx := context.Background()
	ctx = common.WithState(ctx, state)
	ctx = common.WithRuntime(ctx, rt)
	rt.Set("k6", common.Bind(rt, &K6{}, &ctx))

	t.Run("Valid", func(t *testing.T) {
		assert.Equal(t, state.Group, root)
		rt.Set("fn", func() {
			assert.Equal(t, state.Group.Name, "my group")
			assert.Equal(t, state.Group.Parent, root)
		})
		_, err = common.RunString(rt, `k6.group("my group", fn)`)
		assert.NoError(t, err)
		assert.Equal(t, state.Group, root)
	})

	t.Run("Invalid", func(t *testing.T) {
		_, err := common.RunString(rt, `k6.group("::", function() { throw new Error("nooo") })`)
		assert.EqualError(t, err, "GoError: group and check names may not contain '::'")
	})
}

func TestCheck(t *testing.T) {
	rt := goja.New()

	root, err := lib.NewGroup("", nil)
	assert.NoError(t, err)

	baseCtx := common.WithRuntime(context.Background(), rt)

	ctx := new(context.Context)
	*ctx = baseCtx
	rt.Set("k6", common.Bind(rt, &K6{}, ctx))

	t.Run("Object", func(t *testing.T) {
		state := &common.State{Group: root}
		*ctx = common.WithState(baseCtx, state)

		_, err := common.RunString(rt, `k6.check(null, { "check": true })`)
		assert.NoError(t, err)

		if assert.Len(t, state.Samples, 1) {
			assert.NotZero(t, state.Samples[0].Time)
			assert.Equal(t, metrics.Checks, state.Samples[0].Metric)
			assert.Equal(t, float64(1), state.Samples[0].Value)
			assert.Equal(t, map[string]string{
				"group": "",
				"check": "check",
			}, state.Samples[0].Tags)
		}

		t.Run("Invalid", func(t *testing.T) {
			_, err := common.RunString(rt, `k6.check(null, { "::": true })`)
			assert.EqualError(t, err, "GoError: group and check names may not contain '::'")
		})
	})
	t.Run("Array", func(t *testing.T) {
		state := &common.State{Group: root}
		*ctx = common.WithState(baseCtx, state)

		_, err := common.RunString(rt, `k6.check(null, [ true ])`)
		assert.NoError(t, err)

		if assert.Len(t, state.Samples, 1) {
			assert.NotZero(t, state.Samples[0].Time)
			assert.Equal(t, metrics.Checks, state.Samples[0].Metric)
			assert.Equal(t, float64(1), state.Samples[0].Value)
			assert.Equal(t, map[string]string{
				"group": "",
				"check": "0",
			}, state.Samples[0].Tags)
		}
	})
	t.Run("Literal", func(t *testing.T) {
		_, err := common.RunString(rt, `k6.check(null, null)`)
		assert.EqualError(t, err, "TypeError: Cannot convert undefined or null to object")
	})

	t.Run("Throws", func(t *testing.T) {
		_, err := common.RunString(rt, `
		k6.check(null, {
			"a": function() { throw new Error("error A") },
			"b": function() { throw new Error("error B") },
		})
		`)
		assert.EqualError(t, err, "Error: error A at a (<eval>:3:27(6))")
	})

	t.Run("Types", func(t *testing.T) {
		templates := map[string]string{
			"Literal":      `k6.check(null,{"check": %s})`,
			"Callable":     `k6.check(null,{"check": ()=>%s})`,
			"Callable/Arg": `k6.check(%s,{"check":(v)=>v})`,
		}
		testdata := map[string]bool{
			`0`:         false,
			`1`:         true,
			`-1`:        true,
			`""`:        false,
			`"true"`:    true,
			`"false"`:   true,
			`true`:      true,
			`false`:     false,
			`null`:      false,
			`undefined`: false,
		}
		for name, tpl := range templates {
			t.Run(name, func(t *testing.T) {
				for value, succ := range testdata {
					t.Run(value, func(t *testing.T) {
						state := &common.State{Group: root}
						*ctx = common.WithState(baseCtx, state)

						v, err := common.RunString(rt, fmt.Sprintf(tpl, value))
						if assert.NoError(t, err) {
							assert.Equal(t, succ, v.Export())
						}

						if assert.Len(t, state.Samples, 1) {
							assert.NotZero(t, state.Samples[0].Time)
							assert.Equal(t, metrics.Checks, state.Samples[0].Metric)
							if succ {
								assert.Equal(t, float64(1), state.Samples[0].Value)
							} else {
								assert.Equal(t, float64(0), state.Samples[0].Value)
							}
							assert.Equal(t, map[string]string{
								"group": "",
								"check": "check",
							}, state.Samples[0].Tags)
						}
					})
				}
			})
		}
	})

	t.Run("Tags", func(t *testing.T) {
		state := &common.State{Group: root}
		*ctx = common.WithState(baseCtx, state)

		v, err := common.RunString(rt, `k6.check(null, {"check": true}, {a: 1, b: "2"})`)
		if assert.NoError(t, err) {
			assert.Equal(t, true, v.Export())
		}

		if assert.Len(t, state.Samples, 1) {
			assert.NotZero(t, state.Samples[0].Time)
			assert.Equal(t, metrics.Checks, state.Samples[0].Metric)
			assert.Equal(t, float64(1), state.Samples[0].Value)
			assert.Equal(t, map[string]string{
				"group": "",
				"check": "check",
				"a":     "1",
				"b":     "2",
			}, state.Samples[0].Tags)
		}
	})
}
<commit_msg>null isn't a literal, 12345 is a literal<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package k6

import (
	"context"
	"fmt"
	"runtime"
	"testing"
	"time"

	"github.com/dop251/goja"
	"github.com/loadimpact/k6/js/common"
	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/metrics"
	"github.com/stretchr/testify/assert"
)

func TestFail(t *testing.T) {
	rt := goja.New()
	rt.Set("k6", common.Bind(rt, &K6{}, nil))
	_, err := common.RunString(rt, `k6.fail("blah")`)
	assert.EqualError(t, err, "GoError: blah")
}

func TestSleep(t *testing.T) {
	rt := goja.New()
	ctx, cancel := context.WithCancel(context.Background())
	rt.Set("k6", common.Bind(rt, &K6{}, &ctx))

	testdata := map[string]time.Duration{
		"1":   1 * time.Second,
		"1.0": 1 * time.Second,
		"0.5": 500 * time.Millisecond,
	}
	for name, d := range testdata {
		t.Run(name, func(t *testing.T) {
			startTime := time.Now()
			_, err := common.RunString(rt, `k6.sleep(1)`)
			endTime := time.Now()
			assert.NoError(t, err)
			assert.True(t, endTime.Sub(startTime) > d, "did not sleep long enough")
		})
	}

	t.Run("Cancel", func(t *testing.T) {
		dch := make(chan time.Duration)
		go func() {
			startTime := time.Now()
			_, err := common.RunString(rt, `k6.sleep(10)`)
			endTime := time.Now()
			assert.NoError(t, err)
			dch <- endTime.Sub(startTime)
		}()
		runtime.Gosched()
		time.Sleep(1 * time.Second)
		runtime.Gosched()
		cancel()
		runtime.Gosched()
		d := <-dch
		assert.True(t, d > 500*time.Millisecond, "did not sleep long enough")
		assert.True(t, d < 2*time.Second, "slept for too long!!")
	})
}

func TestGroup(t *testing.T) {
	root, err := lib.NewGroup("", nil)
	assert.NoError(t, err)

	rt := goja.New()
	state := &common.State{Group: root}

	ctx := context.Background()
	ctx = common.WithState(ctx, state)
	ctx = common.WithRuntime(ctx, rt)
	rt.Set("k6", common.Bind(rt, &K6{}, &ctx))

	t.Run("Valid", func(t *testing.T) {
		assert.Equal(t, state.Group, root)
		rt.Set("fn", func() {
			assert.Equal(t, state.Group.Name, "my group")
			assert.Equal(t, state.Group.Parent, root)
		})
		_, err = common.RunString(rt, `k6.group("my group", fn)`)
		assert.NoError(t, err)
		assert.Equal(t, state.Group, root)
	})

	t.Run("Invalid", func(t *testing.T) {
		_, err := common.RunString(rt, `k6.group("::", function() { throw new Error("nooo") })`)
		assert.EqualError(t, err, "GoError: group and check names may not contain '::'")
	})
}

func TestCheck(t *testing.T) {
	rt := goja.New()

	root, err := lib.NewGroup("", nil)
	assert.NoError(t, err)

	baseCtx := common.WithRuntime(context.Background(), rt)

	ctx := new(context.Context)
	*ctx = baseCtx
	rt.Set("k6", common.Bind(rt, &K6{}, ctx))

	t.Run("Object", func(t *testing.T) {
		state := &common.State{Group: root}
		*ctx = common.WithState(baseCtx, state)

		_, err := common.RunString(rt, `k6.check(null, { "check": true })`)
		assert.NoError(t, err)

		if assert.Len(t, state.Samples, 1) {
			assert.NotZero(t, state.Samples[0].Time)
			assert.Equal(t, metrics.Checks, state.Samples[0].Metric)
			assert.Equal(t, float64(1), state.Samples[0].Value)
			assert.Equal(t, map[string]string{
				"group": "",
				"check": "check",
			}, state.Samples[0].Tags)
		}

		t.Run("Invalid", func(t *testing.T) {
			_, err := common.RunString(rt, `k6.check(null, { "::": true })`)
			assert.EqualError(t, err, "GoError: group and check names may not contain '::'")
		})
	})
	t.Run("Array", func(t *testing.T) {
		state := &common.State{Group: root}
		*ctx = common.WithState(baseCtx, state)

		_, err := common.RunString(rt, `k6.check(null, [ true ])`)
		assert.NoError(t, err)

		if assert.Len(t, state.Samples, 1) {
			assert.NotZero(t, state.Samples[0].Time)
			assert.Equal(t, metrics.Checks, state.Samples[0].Metric)
			assert.Equal(t, float64(1), state.Samples[0].Value)
			assert.Equal(t, map[string]string{
				"group": "",
				"check": "0",
			}, state.Samples[0].Tags)
		}
	})
	t.Run("Literal", func(t *testing.T) {
		state := &common.State{Group: root}
		*ctx = common.WithState(baseCtx, state)

		_, err := common.RunString(rt, `k6.check(null, 12345)`)
		assert.NoError(t, err)
		assert.Len(t, state.Samples, 0)
	})

	t.Run("Throws", func(t *testing.T) {
		_, err := common.RunString(rt, `
		k6.check(null, {
			"a": function() { throw new Error("error A") },
			"b": function() { throw new Error("error B") },
		})
		`)
		assert.EqualError(t, err, "Error: error A at a (<eval>:3:27(6))")
	})

	t.Run("Types", func(t *testing.T) {
		templates := map[string]string{
			"Literal":      `k6.check(null,{"check": %s})`,
			"Callable":     `k6.check(null,{"check": ()=>%s})`,
			"Callable/Arg": `k6.check(%s,{"check":(v)=>v})`,
		}
		testdata := map[string]bool{
			`0`:         false,
			`1`:         true,
			`-1`:        true,
			`""`:        false,
			`"true"`:    true,
			`"false"`:   true,
			`true`:      true,
			`false`:     false,
			`null`:      false,
			`undefined`: false,
		}
		for name, tpl := range templates {
			t.Run(name, func(t *testing.T) {
				for value, succ := range testdata {
					t.Run(value, func(t *testing.T) {
						state := &common.State{Group: root}
						*ctx = common.WithState(baseCtx, state)

						v, err := common.RunString(rt, fmt.Sprintf(tpl, value))
						if assert.NoError(t, err) {
							assert.Equal(t, succ, v.Export())
						}

						if assert.Len(t, state.Samples, 1) {
							assert.NotZero(t, state.Samples[0].Time)
							assert.Equal(t, metrics.Checks, state.Samples[0].Metric)
							if succ {
								assert.Equal(t, float64(1), state.Samples[0].Value)
							} else {
								assert.Equal(t, float64(0), state.Samples[0].Value)
							}
							assert.Equal(t, map[string]string{
								"group": "",
								"check": "check",
							}, state.Samples[0].Tags)
						}
					})
				}
			})
		}
	})

	t.Run("Tags", func(t *testing.T) {
		state := &common.State{Group: root}
		*ctx = common.WithState(baseCtx, state)

		v, err := common.RunString(rt, `k6.check(null, {"check": true}, {a: 1, b: "2"})`)
		if assert.NoError(t, err) {
			assert.Equal(t, true, v.Export())
		}

		if assert.Len(t, state.Samples, 1) {
			assert.NotZero(t, state.Samples[0].Time)
			assert.Equal(t, metrics.Checks, state.Samples[0].Metric)
			assert.Equal(t, float64(1), state.Samples[0].Value)
			assert.Equal(t, map[string]string{
				"group": "",
				"check": "check",
				"a":     "1",
				"b":     "2",
			}, state.Samples[0].Tags)
		}
	})
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/albertyw/reaction-pics/tumblr"
	// Used for getting tumblr env vars
	_ "github.com/joho/godotenv/autoload"
	newrelic "github.com/newrelic/go-agent"
)

const (
	dataURLPath = "/data.json"
	maxResults  = 20
)

var serverDir = filepath.Join(os.Getenv("ROOT_DIR"), "server")
var staticPath = fmt.Sprintf("%s/static/", serverDir)
var uRLFilePaths = map[string]func() (string, error){}
var board tumblr.Board
var postsMutex sync.RWMutex

// logURL is a closure that logs (to stdout) the url and query of requests
func logURL(
	targetFunc func(http.ResponseWriter, *http.Request),
) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		url := r.URL.String()
		fmt.Println(url)
		targetFunc(w, r)
	}
}

// rewriteFS wraps a static file handler so to rewrite to the static directory
// and the root path is rewritten to index.htm
func rewriteFS(targetFunc func(http.ResponseWriter, *http.Request),
) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		path := strings.TrimPrefix(r.URL.Path, "/static/")
		if path == "/" {
			path = "index.htm"
		}
		r.URL.Path = path
		targetFunc(w, r)
	}
}

// dataURLHandler is an http handler for the dataURLPath response
func dataURLHandler(w http.ResponseWriter, r *http.Request) {
	postsMutex.RLock()
	defer postsMutex.RUnlock()
	board := tumblr.NewBoard(board.Posts)
	html := board.PostsToJSON()
	fmt.Fprintf(w, html)
}

// searchHandler is an http handler to search data for keywords
// It matches the query against post titles and then ranks posts by number of likes
func searchHandler(w http.ResponseWriter, r *http.Request) {
	postsMutex.RLock()
	defer postsMutex.RUnlock()
	query := r.URL.Query().Get("query")
	query = strings.ToLower(query)
	selectedPosts := []tumblr.Post{}
	numPosts := 0
	for _, post := range board.Posts {
		postData := strings.ToLower(post.Title)
		if strings.Contains(postData, query) {
			selectedPosts = append(selectedPosts, post)
			numPosts++
		}
		if numPosts >= maxResults {
			break
		}
	}
	board := tumblr.NewBoard(selectedPosts)
	html := board.PostsToJSON()
	fmt.Fprintf(w, html)
}

func postHandler(w http.ResponseWriter, r *http.Request) {
	pathStrings := strings.Split(r.URL.Path, "/")
	postIDString := pathStrings[2]
	postID, err := strconv.ParseInt(postIDString, 10, 64)
	if err != nil {
		fmt.Println("Cannot parse post id")
		http.NotFound(w, r)
		return
	}
	var post tumblr.Post
	for _, p := range board.Posts {
		if p.ID == postID {
			post = p
		}
	}
	if post == (tumblr.Post{}) {
		fmt.Println("Cannot find post")
		http.NotFound(w, r)
		return
	}
	templateData, err := ioutil.ReadFile(staticPath + "post.htm")
	if err != nil {
		http.Error(w, "Cannot read post template", 500)
		return
	}
	htmlTemplate, err := template.New("page").Parse(string(templateData))
	if err != nil {
		http.Error(w, "Cannot parse post template", 500)
		return
	}
	data := struct {
		Post *tumblr.Post
	}{
		Post: &post,
	}
	err = htmlTemplate.Execute(w, data)
	if err != nil {
		http.Error(w, "Cannot execute post template", 500)
		return
	}
}

// Run starts up the HTTP server
func Run(postChan <-chan tumblr.Post, newrelicApp newrelic.Application) {
	go loadPosts(postChan)
	address := ":" + os.Getenv("PORT")
	fmt.Println("server listening on", address)
	staticFS := rewriteFS(http.FileServer(http.Dir(staticPath)).ServeHTTP)
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/", logURL(staticFS)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, dataURLPath, logURL(dataURLHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/search", logURL(searchHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/post/", logURL(postHandler)))
	http.ListenAndServe(address, nil)
}

func loadPosts(postChan <-chan tumblr.Post) {
	for p := range postChan {
		postsMutex.Lock()
		board.AddPost(p)
		board.SortPostsByLikes()
		postsMutex.Unlock()
	}
}
<commit_msg>Simplify server<commit_after>package server

import (
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/albertyw/reaction-pics/tumblr"
	// Used for getting tumblr env vars
	_ "github.com/joho/godotenv/autoload"
	newrelic "github.com/newrelic/go-agent"
)

const (
	dataURLPath = "/data.json"
	maxResults  = 20
)

var serverDir = filepath.Join(os.Getenv("ROOT_DIR"), "server")
var staticPath = fmt.Sprintf("%s/static/", serverDir)
var uRLFilePaths = map[string]func() (string, error){}
var board = tumblr.NewBoard([]tumblr.Post{})

// logURL is a closure that logs (to stdout) the url and query of requests
func logURL(
	targetFunc func(http.ResponseWriter, *http.Request),
) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		url := r.URL.String()
		fmt.Println(url)
		targetFunc(w, r)
	}
}

// rewriteFS wraps a static file handler so to rewrite to the static directory
// and the root path is rewritten to index.htm
func rewriteFS(targetFunc func(http.ResponseWriter, *http.Request),
) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		path := strings.TrimPrefix(r.URL.Path, "/static/")
		if path == "/" {
			path = "index.htm"
		}
		r.URL.Path = path
		targetFunc(w, r)
	}
}

// dataURLHandler is an http handler for the dataURLPath response
func dataURLHandler(w http.ResponseWriter, r *http.Request) {
	html := board.PostsToJSON()
	fmt.Fprintf(w, html)
}

// searchHandler is an http handler to search data for keywords
// It matches the query against post titles and then ranks posts by number of likes
func searchHandler(w http.ResponseWriter, r *http.Request) {
	query := r.URL.Query().Get("query")
	query = strings.ToLower(query)
	queriedBoard := board.FilterBoard(query, maxResults)
	html := queriedBoard.PostsToJSON()
	fmt.Fprintf(w, html)
}

func postHandler(w http.ResponseWriter, r *http.Request) {
	pathStrings := strings.Split(r.URL.Path, "/")
	postIDString := pathStrings[2]
	postID, err := strconv.ParseInt(postIDString, 10, 64)
	if err != nil {
		fmt.Println("Cannot parse post id")
		http.NotFound(w, r)
		return
	}
	var post tumblr.Post
	for _, p := range board.Posts {
		if p.ID == postID {
			post = p
		}
	}
	if post == (tumblr.Post{}) {
		fmt.Println("Cannot find post")
		http.NotFound(w, r)
		return
	}
	templateData, err := ioutil.ReadFile(staticPath + "post.htm")
	if err != nil {
		http.Error(w, "Cannot read post template", 500)
		return
	}
	htmlTemplate, err := template.New("page").Parse(string(templateData))
	if err != nil {
		http.Error(w, "Cannot parse post template", 500)
		return
	}
	data := struct {
		Post *tumblr.Post
	}{
		Post: &post,
	}
	err = htmlTemplate.Execute(w, data)
	if err != nil {
		http.Error(w, "Cannot execute post template", 500)
		return
	}
}

// Run starts up the HTTP server
func Run(postChan <-chan tumblr.Post, newrelicApp newrelic.Application) {
	go loadPosts(postChan)
	address := ":" + os.Getenv("PORT")
	fmt.Println("server listening on", address)
	staticFS := rewriteFS(http.FileServer(http.Dir(staticPath)).ServeHTTP)
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/", logURL(staticFS)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, dataURLPath, logURL(dataURLHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/search", logURL(searchHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/post/", logURL(postHandler)))
	http.ListenAndServe(address, nil)
}

func loadPosts(postChan <-chan tumblr.Post) {
	for p := range postChan {
		board.AddPost(p)
		board.SortPostsByLikes()
	}
}
<|endoftext|>


<commit_before>package api

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"

	"github.com/ajvb/kala/job"
	"github.com/ajvb/kala/utils/logging"

	"github.com/gorilla/mux"
)

const (
	ApiUrlPrefix = "/api/v1/"
	JobPath      = "job/"
	ApiJobPath   = ApiUrlPrefix + JobPath

	contentType     = "Content-Type"
	jsonContentType = "application/json;charset=UTF-8"
)

var (
	log = logging.GetLogger("api")
)

type KalaStatsResponse struct {
	Stats *job.KalaStats
}

func HandleKalaStatsRequest(w http.ResponseWriter, r *http.Request) {
	resp := &KalaStatsResponse{
		Stats: job.NewKalaStats(),
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}

}

type ListJobStatsResponse struct {
	JobStats []*job.JobStat `json:"job_stats"`
}

func HandleListJobStatsRequest(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]
	j := job.AllJobs.Get(id)

	resp := &ListJobStatsResponse{
		JobStats: j.Stats,
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}

}

type ListJobsResponse struct {
	Jobs map[string]*job.Job `json:"jobs"`
}

// HandleListJobs responds with an array of all Jobs within the server,
// active or disabled.
func HandleListJobsRequest(w http.ResponseWriter, r *http.Request) {
	resp := &ListJobsResponse{
		Jobs: job.AllJobs.GetAll(),
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}

}

type AddJobResponse struct {
	Id string `json:"id"`
}

func unmarshalNewJob(r *http.Request) (*job.Job, error) {
	newJob := &job.Job{}

	body, err := ioutil.ReadAll(io.LimitReader(r.Body, 1048576))
	if err != nil {
		log.Error("Error occured when reading r.Body: %s", err)
		return nil, err
	}
	defer r.Body.Close()

	if err := json.Unmarshal(body, newJob); err != nil {
		log.Error("Error occured when unmarshalling data: %s", err)
		return nil, err
	}

	return newJob, nil
}

// HandleAddJob takes a job object and unmarshals it to a Job type,
// and then throws the job in the schedulers.
func HandleAddJob(w http.ResponseWriter, r *http.Request) {
	newJob, err := unmarshalNewJob(r)
	if err != nil {
		http.Error(w, err.Error(), 400)
		return
	}

	err = newJob.Init()
	if err != nil {
		errStr := "Error occured when initializing the job"
		log.Error(errStr+": %s", err)
		http.Error(w, errStr, 400)
		return
	}
	job.AllJobs.Set(newJob)

	resp := &AddJobResponse{
		Id: newJob.Id,
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}
}

func HandleJobRequest(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]

	if r.Method == "DELETE" {
		handleDeleteJob(w, r, id)
	} else if r.Method == "GET" {
		handleGetJob(w, r, id)
	}
}

func handleDeleteJob(w http.ResponseWriter, r *http.Request, id string) {
	log.Debug("Deleting job: %s", id)

	j := job.AllJobs.Get(id)
	if j == nil {
		w.WriteHeader(http.StatusNotFound)
		return
	}
	j.Delete()

	w.WriteHeader(http.StatusNoContent)
}

type JobResponse struct {
	Job *job.Job `json:"job"`
}

func handleGetJob(w http.ResponseWriter, r *http.Request, id string) {
	j := job.AllJobs.Get(id)
	if j == nil {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	resp := &JobResponse{
		Job: j,
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}
}

func HandleStartJobRequest(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]
	j := job.AllJobs.Get(id)

	j.Run()

	w.WriteHeader(http.StatusNoContent)
}

func SetupApiRoutes(r *mux.Router) {
	// Route for creating a job
	r.HandleFunc(ApiJobPath, HandleAddJob).Methods("POST")
	r.HandleFunc(ApiUrlPrefix+"job", HandleAddJob).Methods("POST")
	// Route for deleting and getting a job
	r.HandleFunc(ApiJobPath+"{id}", HandleJobRequest).Methods("DELETE", "GET")
	r.HandleFunc(ApiJobPath+"{id}/", HandleJobRequest).Methods("DELETE", "GET")
	// Route for getting job stats
	r.HandleFunc(ApiJobPath+"stats/{id}", HandleListJobStatsRequest).Methods("GET")
	r.HandleFunc(ApiJobPath+"stats/{id}/", HandleListJobStatsRequest).Methods("GET")
	// Route for listing all jops
	r.HandleFunc(ApiJobPath, HandleListJobsRequest).Methods("GET")
	r.HandleFunc(ApiUrlPrefix+"job", HandleListJobsRequest).Methods("GET")
	// Route for manually start a job
	r.HandleFunc(ApiJobPath+"start/{id}", HandleStartJobRequest).Methods("POST")
	r.HandleFunc(ApiJobPath+"start/{id}/", HandleStartJobRequest).Methods("POST")
	// Route for getting app-level metrics
	r.HandleFunc(ApiUrlPrefix+"stats", HandleKalaStatsRequest).Methods("GET")
	r.HandleFunc(ApiUrlPrefix+"stats/", HandleKalaStatsRequest).Methods("GET")
}
<commit_msg>Added comments to api.go<commit_after>package api

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"

	"github.com/ajvb/kala/job"
	"github.com/ajvb/kala/utils/logging"

	"github.com/gorilla/mux"
)

const (
	// Base API v1 Path
	ApiUrlPrefix = "/api/v1/"

	JobPath    = "job/"
	ApiJobPath = ApiUrlPrefix + JobPath

	contentType     = "Content-Type"
	jsonContentType = "application/json;charset=UTF-8"
)

var (
	log = logging.GetLogger("api")
)

type KalaStatsResponse struct {
	Stats *job.KalaStats
}

// HandleKalaStatsRequest is the hanlder for getting system-level metrics
// /api/v1/stats
func HandleKalaStatsRequest(w http.ResponseWriter, r *http.Request) {
	resp := &KalaStatsResponse{
		Stats: job.NewKalaStats(),
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}

}

type ListJobStatsResponse struct {
	JobStats []*job.JobStat `json:"job_stats"`
}

// HandleListJobStatsRequest is the handler for getting job-specific stats
// /api/v1/job/stats/{id}
func HandleListJobStatsRequest(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]
	j := job.AllJobs.Get(id)

	resp := &ListJobStatsResponse{
		JobStats: j.Stats,
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}

}

type ListJobsResponse struct {
	Jobs map[string]*job.Job `json:"jobs"`
}

// HandleListJobs responds with an array of all Jobs within the server,
// active or disabled.
func HandleListJobsRequest(w http.ResponseWriter, r *http.Request) {
	resp := &ListJobsResponse{
		Jobs: job.AllJobs.GetAll(),
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}

}

type AddJobResponse struct {
	Id string `json:"id"`
}

func unmarshalNewJob(r *http.Request) (*job.Job, error) {
	newJob := &job.Job{}

	body, err := ioutil.ReadAll(io.LimitReader(r.Body, 1048576))
	if err != nil {
		log.Error("Error occured when reading r.Body: %s", err)
		return nil, err
	}
	defer r.Body.Close()

	if err := json.Unmarshal(body, newJob); err != nil {
		log.Error("Error occured when unmarshalling data: %s", err)
		return nil, err
	}

	return newJob, nil
}

// HandleAddJob takes a job object and unmarshals it to a Job type,
// and then throws the job in the schedulers.
func HandleAddJob(w http.ResponseWriter, r *http.Request) {
	newJob, err := unmarshalNewJob(r)
	if err != nil {
		http.Error(w, err.Error(), 400)
		return
	}

	err = newJob.Init()
	if err != nil {
		errStr := "Error occured when initializing the job"
		log.Error(errStr+": %s", err)
		http.Error(w, errStr, 400)
		return
	}
	job.AllJobs.Set(newJob)

	resp := &AddJobResponse{
		Id: newJob.Id,
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}
}

// HandleJobRequest routes requests to /api/v1/job/{id} to either
// handleDeleteJob if its a DELETE or handleGetJob if its a GET request.
func HandleJobRequest(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]

	if r.Method == "DELETE" {
		handleDeleteJob(w, r, id)
	} else if r.Method == "GET" {
		handleGetJob(w, r, id)
	}
}

func handleDeleteJob(w http.ResponseWriter, r *http.Request, id string) {
	log.Debug("Deleting job: %s", id)

	j := job.AllJobs.Get(id)
	if j == nil {
		w.WriteHeader(http.StatusNotFound)
		return
	}
	j.Delete()

	w.WriteHeader(http.StatusNoContent)
}

type JobResponse struct {
	Job *job.Job `json:"job"`
}

func handleGetJob(w http.ResponseWriter, r *http.Request, id string) {
	j := job.AllJobs.Get(id)
	if j == nil {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	resp := &JobResponse{
		Job: j,
	}

	w.Header().Set(contentType, jsonContentType)
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		log.Error("Error occured when marshalling response: %s", err)
		return
	}
}

// HandleStartJobRequest is the handler for manually starting jobs
// /api/v1/job/start/{id}
func HandleStartJobRequest(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]
	j := job.AllJobs.Get(id)

	j.Run()

	w.WriteHeader(http.StatusNoContent)
}

// SetupApiRoutes is used within main to initialize all of the routes
func SetupApiRoutes(r *mux.Router) {
	// Route for creating a job
	r.HandleFunc(ApiJobPath, HandleAddJob).Methods("POST")
	r.HandleFunc(ApiUrlPrefix+"job", HandleAddJob).Methods("POST")
	// Route for deleting and getting a job
	r.HandleFunc(ApiJobPath+"{id}", HandleJobRequest).Methods("DELETE", "GET")
	r.HandleFunc(ApiJobPath+"{id}/", HandleJobRequest).Methods("DELETE", "GET")
	// Route for getting job stats
	r.HandleFunc(ApiJobPath+"stats/{id}", HandleListJobStatsRequest).Methods("GET")
	r.HandleFunc(ApiJobPath+"stats/{id}/", HandleListJobStatsRequest).Methods("GET")
	// Route for listing all jops
	r.HandleFunc(ApiJobPath, HandleListJobsRequest).Methods("GET")
	r.HandleFunc(ApiUrlPrefix+"job", HandleListJobsRequest).Methods("GET")
	// Route for manually start a job
	r.HandleFunc(ApiJobPath+"start/{id}", HandleStartJobRequest).Methods("POST")
	r.HandleFunc(ApiJobPath+"start/{id}/", HandleStartJobRequest).Methods("POST")
	// Route for getting app-level metrics
	r.HandleFunc(ApiUrlPrefix+"stats", HandleKalaStatsRequest).Methods("GET")
	r.HandleFunc(ApiUrlPrefix+"stats/", HandleKalaStatsRequest).Methods("GET")
}
<|endoftext|>


<commit_before>package ts3

type ServerMethods struct {
	*Client
}

type Channel struct {
	ID                   int    `ms:"cid"`
	ParentID             int    `ms:"pid"`
	ChannelOrder         int    `ms:"channel_order"`
	ChannelName          string `ms:"channel_name"`
	TotalClients         int    `ms:"total_clients"`
	NeededSubscribePower int    `ms:"channel_needed_subscribe_power"`
}

type Instance struct {
	DatabaseVersion             int    `ms:"serverinstance_database_version"`
	FileTransferPort            int    `ms:"serverinstance_filetransfer_port"`
	MaxTotalDownloadBandwidth   uint64 `ms:"serverinstance_max_download_total_bandwidth"`
	MaxTotalUploadBandwidth     uint64 `ms:"serverinstance_max_upload_total_bandwidth"`
	GuestServerQueryGroup       int    `ms:"serverinstance_guest_serverquery_group"`
	ServerQueryFloodCommands    int    `ms:"serverinstance_serverquery_flood_commands"`
	ServerQueryFloodTime        int    `ms:"serverinstance_serverquery_flood_time"`
	ServerQueryBanTime          int    `ms:"serverinstance_serverquery_ban_time"`
	TemplateServerAdminGroup    int    `ms:"serverinstance_template_serveradmin_group"`
	TemplateServerDefaultGroup  int    `ms:"serverinstance_template_serverdefault_group"`
	TemplateChannelAdminGroup   int    `ms:"serverinstance_template_channeladmin_group"`
	TemplateChannelDefaultGroup int    `ms:"serverinstance_template_channeldefault_group"`
}

type ServerConnectionInfo struct {
	FileTransferBandwidthSent     uint64  `ms:"connection_filetransfer_bandwidth_sent"`
	FileTransferBandwidthReceived uint64  `ms:"connection_filetransfer_bandwidth_received"`
	FileTransferTotalSent         uint64  `ms:"connection_filetransfer_bytes_sent_total"`
	FileTransferTotalReceived     uint64  `ms:"connection_filetransfer_bytes_received_total"`
	PacketsSentTotal              uint64  `ms:"connection_packets_sent_total"`
	PacketsReceivedTotal          uint64  `ms:"connection_packets_received_total"`
	BytesSentTotal                uint64  `ms:"connection_bytes_sent_total"`
	BytesReceivedTotal            uint64  `ms:"connection_bytes_received_total"`
	BandwidthSentLastSecond       uint64  `ms:"connection_bandwidth_sent_last_second_total"`
	BandwidthReceivedLastSecond   uint64  `ms:"connection_bandwidth_received_last_second_total"`
	BandwidthSentLastMinute       uint64  `ms:"connection_bandwidth_sent_last_minute_total"`
	BandwidthReceivedLastMinute   uint64  `ms:"connection_bandwidth_received_last_minute_total"`
	ConnectedTime                 uint32  `ms:"connection_connected_time"`
	PacketLossTotalAvg            float32 `ms:"connection_packetloss_total"`
	PingTotalAvg                  float32 `ms:"connection_ping"`
}

// Server represents a TeamSpeak 3 virtual server.
type Server struct {
	AntiFloodPointsNeededCommandBlock      int     `ms:"virtualserver_antiflood_points_needed_command_block"`
	AntiFloodPointsNeededIPBlock           int     `ms:"virtualserver_antiflood_points_needed_ip_block"`
	AntiFloodPointsTickReduce              int     `ms:"virtualserver_antiflood_points_tick_reduce"`
	ChannelsOnline                         int     `ms:"virtualserver_channelsonline"`
	ClientsOnline                          int     `ms:"virtualserver_clientsonline"`
	CodecEncryptionMode                    int     `ms:"virtualserver_codec_encryption_mode"`
	ComplainAutoBanCount                   int     `ms:"virtualserver_complain_autoban_count"`
	ComplainAutoBanTime                    int     `ms:"virtualserver_complain_autoban_time"`
	ComplainRemoveTime                     int     `ms:"virtualserver_complain_remove_time"`
	Created                                int     `ms:"virtualserver_created"`
	DefaultChannelAdminGroup               int     `ms:"virtualserver_default_channel_admin_group"`
	DefaultChannelGroup                    int     `ms:"virtualserver_default_channel_group"`
	DefaultServerGroup                     int     `ms:"virtualserver_default_server_group"`
	HostBannerGFXInterval                  int     `ms:"virtualserver_hostbanner_gfx_interval"`
	HostMessageMode                        int     `ms:"virtualserver_hostmessage_mode"`
	ID                                     int     `ms:"virtualserver_id"`
	MachineID                              string  `ms:"virtualserver_machine_id"`
	MaxClients                             int     `ms:"virtualserver_maxclients"`
	MinClientsInChannelBeforeForcedSilence int     `ms:"virtualserver_min_clients_in_channel_before_forced_silence"`
	NeededIdentitySecurityLevel            int     `ms:"virtualserver_needed_identity_security_level"`
	Port                                   int     `ms:"virtualserver_port"`
	QueryClientsOnline                     int     `ms:"virtualserver_queryclientsonline"`
	Uptime                                 int     `ms:"virtualserver_uptime"` // TODO(steve): use time.Duration
	AskForPrivilegeKey                     bool    `ms:"virtualserver_ask_for_privilegekey"`
	AutoStart                              bool    `ms:"virtualserver_autostart"`
	FlagPassword                           bool    `ms:"virtualserver_flag_password"`
	LogChannel                             bool    `ms:"virtualserver_log_channel"`
	LogClient                              bool    `ms:"virtualserver_log_client"`
	LogFileTransfer                        bool    `ms:"virtualserver_log_filetransfer"`
	LogPermissions                         bool    `ms:"virtualserver_log_permissions"`
	LogQuery                               bool    `ms:"virtualserver_log_client"`
	LogServer                              bool    `ms:"virtualserver_log_server"`
	WebListEnabled                         bool    `ms:"virtualserver_web_list_enabled"`
	PrioritySpeakerDimmModificator         float32 `ms:"virtualserver_priority_speaker_dimm_modificator"`
	BandwidthReceivedLastMinuteTotal       int     `ms:"virtualserver_bandwidth_received_last_minute_total"`
	BandwidthReceivedLastSecondTotal       int     `ms:"virtualserver_bandwidth_received_last_second_total"`
	BandwidthSentLastMinuteTotal           int     `ms:"virtualserver_bandwidth_sent_last_minute_total"`
	BandwidthSentLastSecondTotal           int     `ms:"virtualserver_bandwidth_sent_last_second_total"`
	ChannelTempDeleteDelayDefault          int     `ms:"virtualserver_channel_temp_delete_delay_default"`
	HostBannerMode                         int     `ms:"virtualserver_hostbanner_mode"`
	IconID                                 int     `ms:"virtualserver_icon_id"`
	MinAndroidVersion                      int     `ms:"virtualserver_min_android_version"`
	MinClientVersion                       int     `ms:"virtualserver_min_client_version"`
	MiniOSVersion                          int     `ms:"virtualserver_min_ios_version"`
	ReservedSlots                          int     `ms:"virtualserver_reserved_slots"`
	TotalPing                              float32 `ms:"virtualserver_total_ping"`
	MaxDownloadTotalBandwidth              uint64  `ms:"virtualserver_max_download_total_bandwidth"`
	MaxUploadTotalBandwidth                uint64  `ms:"virtualserver_max_upload_total_bandwidth"`
	MonthBytesDownloaded                   int64   `ms:"virtualserver_MonthBytesDownloaded"`
	MonthBytesUploaded                     int64   `ms:"virtualserver_MonthBytesUploaded"`
	TotalBytesDownloaded                   int64   `ms:"virtualserver_TotalBytesDownloaded"`
	TotalBytesUploaded                     int64   `ms:"virtualserver_TotalBytesUploaded"`
	TotalPacketLossControl                 int64   `ms:"virtualserver_total_packet_loss_control"`
	TotalPacketLossKeepalive               int64   `ms:"virtualserver_total_packet_loss_keepalive"`
	TotalPacketLossSpeech                  int64   `ms:"virtualserver_total_packet_loss_speech"`
	TotalPacketLossTotal                   int64   `ms:"virtualserver_total_packet_loss_total"`
	VirtualServerDownloadQuota             int64   `ms:"virtualserver_virtualserver_download_quota"`
	VirtualServerUploadQuota               int64   `ms:"virtualserver_virtualserver_Upload_quota"`
	FileBase                               string  `ms:"virtualserver_filebase"`
	HostBannerGFXURL                       string  `ms:"virtualserver_hostbanner_gfx_url"`
	HostBannerURL                          string  `ms:"virtualserver_hostbanner_url"`
	HostButtonGFXURL                       string  `ms:"virtualserver_hostbutton_gfx_url"`
	HostButtonToolTip                      string  `ms:"virtualserver_hostbutton_tooltip"`
	HostButtonURL                          string  `ms:"virtualserver_hostbutton_url"`
	HostMessage                            string  `ms:"virtualserver_hostmessage"`
	Name                                   string  `ms:"virtualserver_name"`
	NamePhonetic                           string  `ms:"virtualserver_name_phonetic"`
	Password                               string  `ms:"virtualserver_password"`
	Platform                               string  `ms:"virtualserver_platform"`
	Status                                 string  `ms:"virtualserver_status"`
	UniqueIdentifier                       string  `ms:"virtualserver_unique_identifier"`
	Version                                string  `ms:"virtualserver_version"`
	WelcomeMessage                         string  `ms:"virtualserver_welcomemessage"`
}

// List lists virtual servers.
func (s *ServerMethods) List(options ...string) ([]*Server, error) {
	var servers []*Server
	if _, err := s.ExecCmd(NewCmd("serverlist").WithOptions(options...).WithResponse(&servers)); err != nil {
		return nil, err
	}

	return servers, nil
}

// IDGetByPort returns the database id of the virtual server running on UDP port.
func (s *ServerMethods) IDGetByPort(port uint16) (int, error) {
	r := struct {
		ID int `ms:"server_id"`
	}{}
	_, err := s.ExecCmd(NewCmd("serveridgetbyport").WithArgs(NewArg("virtualserver_port", port)).WithResponse(&r))
	return r.ID, err
}

// Info returns detailed configuration information about the selected server.
func (s *ServerMethods) Info() (*Server, error) {
	r := &Server{}
	if _, err := s.ExecCmd(NewCmd("serverinfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil
}

// InstanceInfo returns detailed information about the selected instance.
func (s *ServerMethods) InstanceInfo() (*Instance, error) {
	r := &Instance{}
	if _, err := s.ExecCmd(NewCmd("instanceinfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil
}

// ServerConnectionInfo returns detailed bandwidth and transfer information about the selected instance.
func (s *ServerMethods) ServerConnectionInfo() (*ServerConnectionInfo, error) {
	r := &ServerConnectionInfo{}
	if _, err := s.ExecCmd(NewCmd("serverrequestconnectioninfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil

}

// Edit changes the selected virtual servers configuration using the given args.
func (s *ServerMethods) Edit(args ...CmdArg) error {
	_, err := s.ExecCmd(NewCmd("serveredit").WithArgs(args...))
	return err
}

// Delete deletes the virtual server specified by id.
// Only virtual server in a stopped state can be deleted.
func (s *ServerMethods) Delete(id int) error {
	_, err := s.ExecCmd(NewCmd("serverdelete").WithArgs(NewArg("sid", id)))
	return err
}

// CreatedServer is the details returned by a server create.
type CreatedServer struct {
	ID    int    `ms:"sid"`
	Port  uint16 `ms:"virtualserver_port"`
	Token string
}

// Create creates a new virtual server using the given properties and returns
// its ID, port and initial administrator privilege key.
// If virtualserver_port arg is not specified, the server will use the first unused
// UDP port.
func (s *ServerMethods) Create(name string, args ...CmdArg) (*CreatedServer, error) {
	r := &CreatedServer{}
	args = append(args, NewArg("virtualserver_name", name))
	if _, err := s.ExecCmd(NewCmd("servercreate").WithArgs(args...).WithResponse(r)); err != nil {
		return nil, err
	}

	return r, nil
}

// Start starts the virtual server specified by id.
func (s *ServerMethods) Start(id int) error {
	_, err := s.ExecCmd(NewCmd("serverstart").WithArgs(NewArg("sid", id)))
	return err
}

// Stop stops the virtual server specified by id.
func (s *ServerMethods) Stop(id int) error {
	_, err := s.ExecCmd(NewCmd("serverstop").WithArgs(NewArg("sid", id)))
	return err
}

// Group represents a virtual server group.
type Group struct {
	ID                int `ms:"sgid"`
	Name              string
	Type              int
	IconID            int
	Saved             bool `ms:"savedb"`
	SortID            int
	NameMode          int
	ModifyPower       int `ms:"n_modifyp"`
	MemberAddPower    int `ms:"n_member_addp"`
	MemberRemovePower int `ms:"n_member_addp"`
}

// GroupList returns a list of available groups for the selected server.
func (s *ServerMethods) GroupList() ([]*Group, error) {
	var groups []*Group
	if _, err := s.ExecCmd(NewCmd("servergrouplist").WithResponse(&groups)); err != nil {
		return nil, err
	}

	return groups, nil
}

// ChannelList returns a list of channels for the selected server.
func (s *ServerMethods) ChannelList() ([]*Channel, error) {
	var channels []*Channel
	if _, err := s.ExecCmd(NewCmd("channellist").WithResponse(&channels)); err != nil {
		return nil, err
	}

	return channels, nil
}

// PrivilegeKey represents a server privilege key.
type PrivilegeKey struct {
	Token       string
	Type        int    `ms:"token_type"`
	ID1         int    `ms:"token_id1"`
	ID2         int    `ms:"token_id2"`
	Created     int    `ms:"token_created"`
	Description string `ms:"token_description"`
}

// PrivilegeKeyList returns a list of available privilege keys for the selected server,
// including their type and group IDs.
func (s *ServerMethods) PrivilegeKeyList() ([]*PrivilegeKey, error) {
	var keys []*PrivilegeKey
	if _, err := s.ExecCmd(NewCmd("privilegekeylist").WithResponse(&keys)); err != nil {
		return nil, err
	}

	return keys, nil
}

// PrivilegeKeyAdd creates a new privilege token to the selected server and returns it.
// If tokentype is set to 0, the ID specified with id1 will be a server group ID.
// Otherwise, id1 is used as a channel group ID and you need to provide a valid channel ID using id2.
func (s *ServerMethods) PrivilegeKeyAdd(ttype, id1, id2 int, options ...CmdArg) (string, error) {
	t := struct {
		Token string
	}{}
	options = append(options, NewArg("tokentype", ttype), NewArg("tokenid1", id1), NewArg("tokenid2", id2))
	_, err := s.ExecCmd(NewCmd("privilegekeylist").WithArgs(options...).WithResponse(&t))
	return t.Token, err
}
<commit_msg>Fixing documentation<commit_after>package ts3

// ServerMethods groups server methods.
type ServerMethods struct {
	*Client
}

// Channel represents a TeamSpeak 3 channel in a virtual server.
type Channel struct {
	ID                   int    `ms:"cid"`
	ParentID             int    `ms:"pid"`
	ChannelOrder         int    `ms:"channel_order"`
	ChannelName          string `ms:"channel_name"`
	TotalClients         int    `ms:"total_clients"`
	NeededSubscribePower int    `ms:"channel_needed_subscribe_power"`
}

// Instance represents basic information for a TeamSpeak 3 instance.
type Instance struct {
	DatabaseVersion             int    `ms:"serverinstance_database_version"`
	FileTransferPort            int    `ms:"serverinstance_filetransfer_port"`
	MaxTotalDownloadBandwidth   uint64 `ms:"serverinstance_max_download_total_bandwidth"`
	MaxTotalUploadBandwidth     uint64 `ms:"serverinstance_max_upload_total_bandwidth"`
	GuestServerQueryGroup       int    `ms:"serverinstance_guest_serverquery_group"`
	ServerQueryFloodCommands    int    `ms:"serverinstance_serverquery_flood_commands"`
	ServerQueryFloodTime        int    `ms:"serverinstance_serverquery_flood_time"`
	ServerQueryBanTime          int    `ms:"serverinstance_serverquery_ban_time"`
	TemplateServerAdminGroup    int    `ms:"serverinstance_template_serveradmin_group"`
	TemplateServerDefaultGroup  int    `ms:"serverinstance_template_serverdefault_group"`
	TemplateChannelAdminGroup   int    `ms:"serverinstance_template_channeladmin_group"`
	TemplateChannelDefaultGroup int    `ms:"serverinstance_template_channeldefault_group"`
}

// ServerConnectionInfo represents the connection info for a TeamSpeak 3 instance.
type ServerConnectionInfo struct {
	FileTransferBandwidthSent     uint64  `ms:"connection_filetransfer_bandwidth_sent"`
	FileTransferBandwidthReceived uint64  `ms:"connection_filetransfer_bandwidth_received"`
	FileTransferTotalSent         uint64  `ms:"connection_filetransfer_bytes_sent_total"`
	FileTransferTotalReceived     uint64  `ms:"connection_filetransfer_bytes_received_total"`
	PacketsSentTotal              uint64  `ms:"connection_packets_sent_total"`
	PacketsReceivedTotal          uint64  `ms:"connection_packets_received_total"`
	BytesSentTotal                uint64  `ms:"connection_bytes_sent_total"`
	BytesReceivedTotal            uint64  `ms:"connection_bytes_received_total"`
	BandwidthSentLastSecond       uint64  `ms:"connection_bandwidth_sent_last_second_total"`
	BandwidthReceivedLastSecond   uint64  `ms:"connection_bandwidth_received_last_second_total"`
	BandwidthSentLastMinute       uint64  `ms:"connection_bandwidth_sent_last_minute_total"`
	BandwidthReceivedLastMinute   uint64  `ms:"connection_bandwidth_received_last_minute_total"`
	ConnectedTime                 uint32  `ms:"connection_connected_time"`
	PacketLossTotalAvg            float32 `ms:"connection_packetloss_total"`
	PingTotalAvg                  float32 `ms:"connection_ping"`
}

// Server represents a TeamSpeak 3 virtual server.
type Server struct {
	AntiFloodPointsNeededCommandBlock      int     `ms:"virtualserver_antiflood_points_needed_command_block"`
	AntiFloodPointsNeededIPBlock           int     `ms:"virtualserver_antiflood_points_needed_ip_block"`
	AntiFloodPointsTickReduce              int     `ms:"virtualserver_antiflood_points_tick_reduce"`
	ChannelsOnline                         int     `ms:"virtualserver_channelsonline"`
	ClientsOnline                          int     `ms:"virtualserver_clientsonline"`
	CodecEncryptionMode                    int     `ms:"virtualserver_codec_encryption_mode"`
	ComplainAutoBanCount                   int     `ms:"virtualserver_complain_autoban_count"`
	ComplainAutoBanTime                    int     `ms:"virtualserver_complain_autoban_time"`
	ComplainRemoveTime                     int     `ms:"virtualserver_complain_remove_time"`
	Created                                int     `ms:"virtualserver_created"`
	DefaultChannelAdminGroup               int     `ms:"virtualserver_default_channel_admin_group"`
	DefaultChannelGroup                    int     `ms:"virtualserver_default_channel_group"`
	DefaultServerGroup                     int     `ms:"virtualserver_default_server_group"`
	HostBannerGFXInterval                  int     `ms:"virtualserver_hostbanner_gfx_interval"`
	HostMessageMode                        int     `ms:"virtualserver_hostmessage_mode"`
	ID                                     int     `ms:"virtualserver_id"`
	MachineID                              string  `ms:"virtualserver_machine_id"`
	MaxClients                             int     `ms:"virtualserver_maxclients"`
	MinClientsInChannelBeforeForcedSilence int     `ms:"virtualserver_min_clients_in_channel_before_forced_silence"`
	NeededIdentitySecurityLevel            int     `ms:"virtualserver_needed_identity_security_level"`
	Port                                   int     `ms:"virtualserver_port"`
	QueryClientsOnline                     int     `ms:"virtualserver_queryclientsonline"`
	Uptime                                 int     `ms:"virtualserver_uptime"` // TODO(steve): use time.Duration
	AskForPrivilegeKey                     bool    `ms:"virtualserver_ask_for_privilegekey"`
	AutoStart                              bool    `ms:"virtualserver_autostart"`
	FlagPassword                           bool    `ms:"virtualserver_flag_password"`
	LogChannel                             bool    `ms:"virtualserver_log_channel"`
	LogClient                              bool    `ms:"virtualserver_log_client"`
	LogFileTransfer                        bool    `ms:"virtualserver_log_filetransfer"`
	LogPermissions                         bool    `ms:"virtualserver_log_permissions"`
	LogQuery                               bool    `ms:"virtualserver_log_client"`
	LogServer                              bool    `ms:"virtualserver_log_server"`
	WebListEnabled                         bool    `ms:"virtualserver_web_list_enabled"`
	PrioritySpeakerDimmModificator         float32 `ms:"virtualserver_priority_speaker_dimm_modificator"`
	BandwidthReceivedLastMinuteTotal       int     `ms:"virtualserver_bandwidth_received_last_minute_total"`
	BandwidthReceivedLastSecondTotal       int     `ms:"virtualserver_bandwidth_received_last_second_total"`
	BandwidthSentLastMinuteTotal           int     `ms:"virtualserver_bandwidth_sent_last_minute_total"`
	BandwidthSentLastSecondTotal           int     `ms:"virtualserver_bandwidth_sent_last_second_total"`
	ChannelTempDeleteDelayDefault          int     `ms:"virtualserver_channel_temp_delete_delay_default"`
	HostBannerMode                         int     `ms:"virtualserver_hostbanner_mode"`
	IconID                                 int     `ms:"virtualserver_icon_id"`
	MinAndroidVersion                      int     `ms:"virtualserver_min_android_version"`
	MinClientVersion                       int     `ms:"virtualserver_min_client_version"`
	MiniOSVersion                          int     `ms:"virtualserver_min_ios_version"`
	ReservedSlots                          int     `ms:"virtualserver_reserved_slots"`
	TotalPing                              float32 `ms:"virtualserver_total_ping"`
	MaxDownloadTotalBandwidth              uint64  `ms:"virtualserver_max_download_total_bandwidth"`
	MaxUploadTotalBandwidth                uint64  `ms:"virtualserver_max_upload_total_bandwidth"`
	MonthBytesDownloaded                   int64   `ms:"virtualserver_MonthBytesDownloaded"`
	MonthBytesUploaded                     int64   `ms:"virtualserver_MonthBytesUploaded"`
	TotalBytesDownloaded                   int64   `ms:"virtualserver_TotalBytesDownloaded"`
	TotalBytesUploaded                     int64   `ms:"virtualserver_TotalBytesUploaded"`
	TotalPacketLossControl                 int64   `ms:"virtualserver_total_packet_loss_control"`
	TotalPacketLossKeepalive               int64   `ms:"virtualserver_total_packet_loss_keepalive"`
	TotalPacketLossSpeech                  int64   `ms:"virtualserver_total_packet_loss_speech"`
	TotalPacketLossTotal                   int64   `ms:"virtualserver_total_packet_loss_total"`
	VirtualServerDownloadQuota             int64   `ms:"virtualserver_virtualserver_download_quota"`
	VirtualServerUploadQuota               int64   `ms:"virtualserver_virtualserver_Upload_quota"`
	FileBase                               string  `ms:"virtualserver_filebase"`
	HostBannerGFXURL                       string  `ms:"virtualserver_hostbanner_gfx_url"`
	HostBannerURL                          string  `ms:"virtualserver_hostbanner_url"`
	HostButtonGFXURL                       string  `ms:"virtualserver_hostbutton_gfx_url"`
	HostButtonToolTip                      string  `ms:"virtualserver_hostbutton_tooltip"`
	HostButtonURL                          string  `ms:"virtualserver_hostbutton_url"`
	HostMessage                            string  `ms:"virtualserver_hostmessage"`
	Name                                   string  `ms:"virtualserver_name"`
	NamePhonetic                           string  `ms:"virtualserver_name_phonetic"`
	Password                               string  `ms:"virtualserver_password"`
	Platform                               string  `ms:"virtualserver_platform"`
	Status                                 string  `ms:"virtualserver_status"`
	UniqueIdentifier                       string  `ms:"virtualserver_unique_identifier"`
	Version                                string  `ms:"virtualserver_version"`
	WelcomeMessage                         string  `ms:"virtualserver_welcomemessage"`
}

// List lists virtual servers.
func (s *ServerMethods) List(options ...string) ([]*Server, error) {
	var servers []*Server
	if _, err := s.ExecCmd(NewCmd("serverlist").WithOptions(options...).WithResponse(&servers)); err != nil {
		return nil, err
	}

	return servers, nil
}

// IDGetByPort returns the database id of the virtual server running on UDP port.
func (s *ServerMethods) IDGetByPort(port uint16) (int, error) {
	r := struct {
		ID int `ms:"server_id"`
	}{}
	_, err := s.ExecCmd(NewCmd("serveridgetbyport").WithArgs(NewArg("virtualserver_port", port)).WithResponse(&r))
	return r.ID, err
}

// Info returns detailed configuration information about the selected server.
func (s *ServerMethods) Info() (*Server, error) {
	r := &Server{}
	if _, err := s.ExecCmd(NewCmd("serverinfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil
}

// InstanceInfo returns detailed information about the selected instance.
func (s *ServerMethods) InstanceInfo() (*Instance, error) {
	r := &Instance{}
	if _, err := s.ExecCmd(NewCmd("instanceinfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil
}

// ServerConnectionInfo returns detailed bandwidth and transfer information about the selected instance.
func (s *ServerMethods) ServerConnectionInfo() (*ServerConnectionInfo, error) {
	r := &ServerConnectionInfo{}
	if _, err := s.ExecCmd(NewCmd("serverrequestconnectioninfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil

}

// Edit changes the selected virtual servers configuration using the given args.
func (s *ServerMethods) Edit(args ...CmdArg) error {
	_, err := s.ExecCmd(NewCmd("serveredit").WithArgs(args...))
	return err
}

// Delete deletes the virtual server specified by id.
// Only virtual server in a stopped state can be deleted.
func (s *ServerMethods) Delete(id int) error {
	_, err := s.ExecCmd(NewCmd("serverdelete").WithArgs(NewArg("sid", id)))
	return err
}

// CreatedServer is the details returned by a server create.
type CreatedServer struct {
	ID    int    `ms:"sid"`
	Port  uint16 `ms:"virtualserver_port"`
	Token string
}

// Create creates a new virtual server using the given properties and returns
// its ID, port and initial administrator privilege key.
// If virtualserver_port arg is not specified, the server will use the first unused
// UDP port.
func (s *ServerMethods) Create(name string, args ...CmdArg) (*CreatedServer, error) {
	r := &CreatedServer{}
	args = append(args, NewArg("virtualserver_name", name))
	if _, err := s.ExecCmd(NewCmd("servercreate").WithArgs(args...).WithResponse(r)); err != nil {
		return nil, err
	}

	return r, nil
}

// Start starts the virtual server specified by id.
func (s *ServerMethods) Start(id int) error {
	_, err := s.ExecCmd(NewCmd("serverstart").WithArgs(NewArg("sid", id)))
	return err
}

// Stop stops the virtual server specified by id.
func (s *ServerMethods) Stop(id int) error {
	_, err := s.ExecCmd(NewCmd("serverstop").WithArgs(NewArg("sid", id)))
	return err
}

// Group represents a virtual server group.
type Group struct {
	ID                int `ms:"sgid"`
	Name              string
	Type              int
	IconID            int
	Saved             bool `ms:"savedb"`
	SortID            int
	NameMode          int
	ModifyPower       int `ms:"n_modifyp"`
	MemberAddPower    int `ms:"n_member_addp"`
	MemberRemovePower int `ms:"n_member_addp"`
}

// GroupList returns a list of available groups for the selected server.
func (s *ServerMethods) GroupList() ([]*Group, error) {
	var groups []*Group
	if _, err := s.ExecCmd(NewCmd("servergrouplist").WithResponse(&groups)); err != nil {
		return nil, err
	}

	return groups, nil
}

// ChannelList returns a list of channels for the selected server.
func (s *ServerMethods) ChannelList() ([]*Channel, error) {
	var channels []*Channel
	if _, err := s.ExecCmd(NewCmd("channellist").WithResponse(&channels)); err != nil {
		return nil, err
	}

	return channels, nil
}

// PrivilegeKey represents a server privilege key.
type PrivilegeKey struct {
	Token       string
	Type        int    `ms:"token_type"`
	ID1         int    `ms:"token_id1"`
	ID2         int    `ms:"token_id2"`
	Created     int    `ms:"token_created"`
	Description string `ms:"token_description"`
}

// PrivilegeKeyList returns a list of available privilege keys for the selected server,
// including their type and group IDs.
func (s *ServerMethods) PrivilegeKeyList() ([]*PrivilegeKey, error) {
	var keys []*PrivilegeKey
	if _, err := s.ExecCmd(NewCmd("privilegekeylist").WithResponse(&keys)); err != nil {
		return nil, err
	}

	return keys, nil
}

// PrivilegeKeyAdd creates a new privilege token to the selected server and returns it.
// If tokentype is set to 0, the ID specified with id1 will be a server group ID.
// Otherwise, id1 is used as a channel group ID and you need to provide a valid channel ID using id2.
func (s *ServerMethods) PrivilegeKeyAdd(ttype, id1, id2 int, options ...CmdArg) (string, error) {
	t := struct {
		Token string
	}{}
	options = append(options, NewArg("tokentype", ttype), NewArg("tokenid1", id1), NewArg("tokenid2", id2))
	_, err := s.ExecCmd(NewCmd("privilegekeylist").WithArgs(options...).WithResponse(&t))
	return t.Token, err
}
<|endoftext|>