<commit_before>package hoverfly

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	// static assets
	_ "github.com/SpectoLabs/hoverfly/statik"
	"github.com/rakyll/statik/fs"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/go-zoo/bone"
	"github.com/gorilla/websocket"
	"github.com/meatballhat/negroni-logrus"

	// auth
	"github.com/SpectoLabs/hoverfly/authentication"
	"github.com/SpectoLabs/hoverfly/authentication/controllers"
)

// recordedRequests struct encapsulates payload data
type recordedRequests struct {
	Data []Payload `json:"data"`
}

type recordsCount struct {
	Count int `json:"count"`
}

type statsResponse struct {
	Stats        Stats `json:"stats"`
	RecordsCount int   `json:"recordsCount"`
}

type stateRequest struct {
	Mode        string `json:"mode"`
	Destination string `json:"destination"`
}

type messageResponse struct {
	Message string `json:"message"`
}

// StartAdminInterface - starts admin interface web server
func (d *DBClient) StartAdminInterface() {
	go func() {
		// starting admin interface
		mux := getBoneRouter(*d)
		n := negroni.Classic()

		logLevel := log.ErrorLevel

		if d.Cfg.Verbose {
			logLevel = log.DebugLevel
		}

		n.Use(negronilogrus.NewCustomMiddleware(logLevel, &log.JSONFormatter{}, "admin"))
		n.UseHandler(mux)

		// admin interface starting message
		log.WithFields(log.Fields{
			"AdminPort": d.Cfg.AdminPort,
		}).Info("Admin interface is starting...")

		n.Run(fmt.Sprintf(":%s", d.Cfg.AdminPort))
	}()
}

// getBoneRouter returns mux for admin interface
func getBoneRouter(d DBClient) *bone.Mux {
	mux := bone.New()

	// getting auth controllers and middleware
	ac := controllers.GetNewAuthenticationController(d.AB, d.Cfg.SecretKey, d.Cfg.JWTExpirationDelta)
	am := authentication.GetNewAuthenticationMiddleware(d.AB, d.Cfg.SecretKey, d.Cfg.JWTExpirationDelta)

	mux.Post("/token-auth", http.HandlerFunc(ac.Login))
	mux.Get("/refresh-token-auth", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.RefreshToken),
	))
	mux.Get("/logout", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.Logout),
	))

	mux.Get("/users", http.HandlerFunc(ac.GetAllUsersHandler))
	// TODO: add users delete/add functionality

	mux.Get("/records", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(d.AllRecordsHandler),
	))
	mux.Delete("/records", http.HandlerFunc(d.DeleteAllRecordsHandler))
	mux.Post("/records", http.HandlerFunc(d.ImportRecordsHandler))

	mux.Get("/count", http.HandlerFunc(d.RecordsCount))
	mux.Get("/stats", http.HandlerFunc(d.StatsHandler))
	mux.Get("/statsws", http.HandlerFunc(d.StatsWSHandler))

	mux.Get("/state", http.HandlerFunc(d.CurrentStateHandler))
	mux.Post("/state", http.HandlerFunc(d.StateHandler))

	if d.Cfg.Development {
		// since hoverfly is not started from cmd/hoverfly/hoverfly
		// we have to target to that directory
		log.Warn("Hoverfly is serving files from /static/dist instead of statik binary!")
		mux.Handle("/*", http.FileServer(http.Dir("../../static/dist")))
	} else {
		// preparing static assets for embedded admin
		statikFS, err := fs.New()

		if err != nil {
			log.WithFields(log.Fields{
				"Error": err.Error(),
			}).Error("Failed to load statikFS, admin UI might not work :(")
		}

		mux.Handle("/*", http.FileServer(statikFS))
	}

	return mux
}

// AllRecordsHandler returns JSON content type http response
func (d *DBClient) AllRecordsHandler(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {
	records, err := d.Cache.GetAllRequests()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordedRequests
		response.Data = records
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// RecordsCount returns number of captured requests as a JSON payload
func (d *DBClient) RecordsCount(w http.ResponseWriter, req *http.Request) {
	count, err := d.Cache.RecordsCount()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordsCount
		response.Count = count
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// StatsHandler - returns current stats about Hoverfly (request counts, record count)
func (d *DBClient) StatsHandler(w http.ResponseWriter, req *http.Request) {
	stats := d.Counter.Flush()

	count, err := d.Cache.RecordsCount()

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	var sr statsResponse
	sr.Stats = stats
	sr.RecordsCount = count

	w.Header().Set("Content-Type", "application/json")

	b, err := json.Marshal(sr)

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write(b)
		return
	}

}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// StatsWSHandler - returns current stats about Hoverfly (request counts, record count) through the websocket
func (d *DBClient) StatsWSHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			return
		}
		log.WithFields(log.Fields{
			"message": string(p),
		}).Info("Got message...")

		for _ = range time.Tick(1 * time.Second) {

			count, err := d.Cache.RecordsCount()

			if err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("got error while trying to get records count")
				return
			}

			stats := d.Counter.Flush()

			var sr statsResponse
			sr.Stats = stats
			sr.RecordsCount = count

			b, err := json.Marshal(sr)

			if err = conn.WriteMessage(messageType, b); err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Debug("Got error when writing message...")
				return
			}
		}

	}

}

// ImportRecordsHandler - accepts JSON payload and saves it to cache
func (d *DBClient) ImportRecordsHandler(w http.ResponseWriter, req *http.Request) {

	var requests recordedRequests

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	var response messageResponse

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		response.Message = "Bad request. Nothing to import!"
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &requests)

	if err != nil {
		w.WriteHeader(422) // can't process this entity
		return
	}

	err = d.ImportPayloads(requests.Data)

	if err != nil {
		response.Message = err.Error()
		w.WriteHeader(400)
	} else {
		response.Message = fmt.Sprintf("%d payloads import complete.", len(requests.Data))
	}

	b, err := json.Marshal(response)
	w.Write(b)

}

// DeleteAllRecordsHandler - deletes all captured requests
func (d *DBClient) DeleteAllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	err := d.Cache.DeleteData()

	var en Entry
	en.ActionType = ActionTypeWipeDB
	en.Message = "wipe"
	en.Time = time.Now()

	if err := d.Hooks.Fire(ActionTypeWipeDB, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeWipeDB,
		}).Error("failed to fire hook")
	}

	w.Header().Set("Content-Type", "application/json")

	var response messageResponse
	if err != nil {
		if err.Error() == "bucket not found" {
			response.Message = fmt.Sprintf("No records found")
			w.WriteHeader(200)
		} else {
			response.Message = fmt.Sprintf("Something went wrong: %s", err.Error())
			w.WriteHeader(500)
		}
	} else {
		response.Message = "Proxy cache deleted successfuly"
		w.WriteHeader(200)
	}
	b, err := json.Marshal(response)

	w.Write(b)
	return
}

// CurrentStateHandler returns current state
func (d *DBClient) CurrentStateHandler(w http.ResponseWriter, req *http.Request) {
	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination

	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)
}

// StateHandler handles current proxy state
func (d *DBClient) StateHandler(w http.ResponseWriter, r *http.Request) {
	var sr stateRequest

	// this is mainly for testing, since when you create
	if r.Body == nil {
		r.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("")))
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &sr)

	if err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(400) // can't process this entity
		return
	}

	availableModes := map[string]bool{
		"virtualize": true,
		"capture":    true,
		"modify":     true,
		"synthesize": true,
	}

	if !availableModes[sr.Mode] {
		log.WithFields(log.Fields{
			"suppliedMode": sr.Mode,
		}).Error("Wrong mode found, can't change state")
		http.Error(w, "Bad mode supplied, available modes: virtualize, capture, modify, synthesize.", 400)
		return
	}

	log.WithFields(log.Fields{
		"newState": sr.Mode,
		"body":     string(body),
	}).Info("Handling state change request!")

	// setting new state
	d.Cfg.SetMode(sr.Mode)

	var en Entry
	en.ActionType = ActionTypeConfigurationChanged
	en.Message = "changed"
	en.Time = time.Now()
	en.Data = []byte("sr.Mode")

	if err := d.Hooks.Fire(ActionTypeConfigurationChanged, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeConfigurationChanged,
		}).Error("failed to fire hook")
	}

	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination
	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)

}
<commit_msg>passing auth enabled variable to middleware<commit_after>package hoverfly

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	// static assets
	_ "github.com/SpectoLabs/hoverfly/statik"
	"github.com/rakyll/statik/fs"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/go-zoo/bone"
	"github.com/gorilla/websocket"
	"github.com/meatballhat/negroni-logrus"

	// auth
	"github.com/SpectoLabs/hoverfly/authentication"
	"github.com/SpectoLabs/hoverfly/authentication/controllers"
)

// recordedRequests struct encapsulates payload data
type recordedRequests struct {
	Data []Payload `json:"data"`
}

type recordsCount struct {
	Count int `json:"count"`
}

type statsResponse struct {
	Stats        Stats `json:"stats"`
	RecordsCount int   `json:"recordsCount"`
}

type stateRequest struct {
	Mode        string `json:"mode"`
	Destination string `json:"destination"`
}

type messageResponse struct {
	Message string `json:"message"`
}

// StartAdminInterface - starts admin interface web server
func (d *DBClient) StartAdminInterface() {
	go func() {
		// starting admin interface
		mux := getBoneRouter(*d)
		n := negroni.Classic()

		logLevel := log.ErrorLevel

		if d.Cfg.Verbose {
			logLevel = log.DebugLevel
		}

		n.Use(negronilogrus.NewCustomMiddleware(logLevel, &log.JSONFormatter{}, "admin"))
		n.UseHandler(mux)

		// admin interface starting message
		log.WithFields(log.Fields{
			"AdminPort": d.Cfg.AdminPort,
		}).Info("Admin interface is starting...")

		n.Run(fmt.Sprintf(":%s", d.Cfg.AdminPort))
	}()
}

// getBoneRouter returns mux for admin interface
func getBoneRouter(d DBClient) *bone.Mux {
	mux := bone.New()

	// getting auth controllers and middleware
	ac := controllers.GetNewAuthenticationController(d.AB, d.Cfg.SecretKey, d.Cfg.JWTExpirationDelta)
	am := authentication.GetNewAuthenticationMiddleware(d.AB,
		d.Cfg.SecretKey,
		d.Cfg.JWTExpirationDelta,
		d.Cfg.AuthEnabled)

	mux.Post("/token-auth", http.HandlerFunc(ac.Login))
	mux.Get("/refresh-token-auth", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.RefreshToken),
	))
	mux.Get("/logout", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.Logout),
	))

	mux.Get("/users", http.HandlerFunc(ac.GetAllUsersHandler))
	// TODO: add users delete/add functionality

	mux.Get("/records", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(d.AllRecordsHandler),
	))
	mux.Delete("/records", http.HandlerFunc(d.DeleteAllRecordsHandler))
	mux.Post("/records", http.HandlerFunc(d.ImportRecordsHandler))

	mux.Get("/count", http.HandlerFunc(d.RecordsCount))
	mux.Get("/stats", http.HandlerFunc(d.StatsHandler))
	mux.Get("/statsws", http.HandlerFunc(d.StatsWSHandler))

	mux.Get("/state", http.HandlerFunc(d.CurrentStateHandler))
	mux.Post("/state", http.HandlerFunc(d.StateHandler))

	if d.Cfg.Development {
		// since hoverfly is not started from cmd/hoverfly/hoverfly
		// we have to target to that directory
		log.Warn("Hoverfly is serving files from /static/dist instead of statik binary!")
		mux.Handle("/*", http.FileServer(http.Dir("../../static/dist")))
	} else {
		// preparing static assets for embedded admin
		statikFS, err := fs.New()

		if err != nil {
			log.WithFields(log.Fields{
				"Error": err.Error(),
			}).Error("Failed to load statikFS, admin UI might not work :(")
		}

		mux.Handle("/*", http.FileServer(statikFS))
	}
	return mux
}

// AllRecordsHandler returns JSON content type http response
func (d *DBClient) AllRecordsHandler(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {
	records, err := d.Cache.GetAllRequests()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordedRequests
		response.Data = records
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// RecordsCount returns number of captured requests as a JSON payload
func (d *DBClient) RecordsCount(w http.ResponseWriter, req *http.Request) {
	count, err := d.Cache.RecordsCount()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordsCount
		response.Count = count
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// StatsHandler - returns current stats about Hoverfly (request counts, record count)
func (d *DBClient) StatsHandler(w http.ResponseWriter, req *http.Request) {
	stats := d.Counter.Flush()

	count, err := d.Cache.RecordsCount()

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	var sr statsResponse
	sr.Stats = stats
	sr.RecordsCount = count

	w.Header().Set("Content-Type", "application/json")

	b, err := json.Marshal(sr)

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write(b)
		return
	}

}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// StatsWSHandler - returns current stats about Hoverfly (request counts, record count) through the websocket
func (d *DBClient) StatsWSHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			return
		}
		log.WithFields(log.Fields{
			"message": string(p),
		}).Info("Got message...")

		for _ = range time.Tick(1 * time.Second) {

			count, err := d.Cache.RecordsCount()

			if err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("got error while trying to get records count")
				return
			}

			stats := d.Counter.Flush()

			var sr statsResponse
			sr.Stats = stats
			sr.RecordsCount = count

			b, err := json.Marshal(sr)

			if err = conn.WriteMessage(messageType, b); err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Debug("Got error when writing message...")
				return
			}
		}

	}

}

// ImportRecordsHandler - accepts JSON payload and saves it to cache
func (d *DBClient) ImportRecordsHandler(w http.ResponseWriter, req *http.Request) {

	var requests recordedRequests

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	var response messageResponse

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		response.Message = "Bad request. Nothing to import!"
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &requests)

	if err != nil {
		w.WriteHeader(422) // can't process this entity
		return
	}

	err = d.ImportPayloads(requests.Data)

	if err != nil {
		response.Message = err.Error()
		w.WriteHeader(400)
	} else {
		response.Message = fmt.Sprintf("%d payloads import complete.", len(requests.Data))
	}

	b, err := json.Marshal(response)
	w.Write(b)

}

// DeleteAllRecordsHandler - deletes all captured requests
func (d *DBClient) DeleteAllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	err := d.Cache.DeleteData()

	var en Entry
	en.ActionType = ActionTypeWipeDB
	en.Message = "wipe"
	en.Time = time.Now()

	if err := d.Hooks.Fire(ActionTypeWipeDB, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeWipeDB,
		}).Error("failed to fire hook")
	}

	w.Header().Set("Content-Type", "application/json")

	var response messageResponse
	if err != nil {
		if err.Error() == "bucket not found" {
			response.Message = fmt.Sprintf("No records found")
			w.WriteHeader(200)
		} else {
			response.Message = fmt.Sprintf("Something went wrong: %s", err.Error())
			w.WriteHeader(500)
		}
	} else {
		response.Message = "Proxy cache deleted successfuly"
		w.WriteHeader(200)
	}
	b, err := json.Marshal(response)

	w.Write(b)
	return
}

// CurrentStateHandler returns current state
func (d *DBClient) CurrentStateHandler(w http.ResponseWriter, req *http.Request) {
	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination

	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)
}

// StateHandler handles current proxy state
func (d *DBClient) StateHandler(w http.ResponseWriter, r *http.Request) {
	var sr stateRequest

	// this is mainly for testing, since when you create
	if r.Body == nil {
		r.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("")))
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &sr)

	if err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(400) // can't process this entity
		return
	}

	availableModes := map[string]bool{
		"virtualize": true,
		"capture":    true,
		"modify":     true,
		"synthesize": true,
	}

	if !availableModes[sr.Mode] {
		log.WithFields(log.Fields{
			"suppliedMode": sr.Mode,
		}).Error("Wrong mode found, can't change state")
		http.Error(w, "Bad mode supplied, available modes: virtualize, capture, modify, synthesize.", 400)
		return
	}

	log.WithFields(log.Fields{
		"newState": sr.Mode,
		"body":     string(body),
	}).Info("Handling state change request!")

	// setting new state
	d.Cfg.SetMode(sr.Mode)

	var en Entry
	en.ActionType = ActionTypeConfigurationChanged
	en.Message = "changed"
	en.Time = time.Now()
	en.Data = []byte("sr.Mode")

	if err := d.Hooks.Fire(ActionTypeConfigurationChanged, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeConfigurationChanged,
		}).Error("failed to fire hook")
	}

	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination
	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"sync"
)

type Conf struct {
	data    map[string]string
	dataMux sync.RWMutex
}

func (c *Conf) Get(k string) string {
	// Locking in base function
	return c.GetOrDefault(k, "")
}

func (c *Conf) Set(k string, v string) {
	c.dataMux.Lock()
	defer c.dataMux.Unlock()
	c.data[k] = v
}

func (c *Conf) Save() bool {
	c.dataMux.RLock()
	defer c.dataMux.RUnlock()
	b, je := json.Marshal(c.data)
	if je != nil {
		log.Printf("Failed saving conf: %s", je)
		return false
	}
	we := ioutil.WriteFile(confPath, b, 0600)
	if we != nil {
		log.Printf("Failed saving conf: %s", we)
		return false
	}
	return true
}

func (c *Conf) GetNotEmpty(k string) string {
	// Locking in base function
	val := c.GetOrDefault(k, "")
	if len(val) < 1 {
		panic(fmt.Sprintf("Value %s empty", k))
	}
	return val
}

func (c *Conf) GetOrDefault(k string, d string) string {
	c.dataMux.RLock()
	defer c.dataMux.RUnlock()
	if len(c.data[k]) == 0 {
		return d
	}
	return c.data[k]
}

func newConf(path string) *Conf {
	c := &Conf{}
	if len(path) > 0 {
		// Load file
		b, e := ioutil.ReadFile(path)
		if e != nil {
			log.Fatal(fmt.Sprintf("Failed to load conf: %s", e))
		}

		// Parse JSON
		var data map[string]string
		je := json.Unmarshal(b, &data)
		if je != nil {
			log.Fatal(fmt.Sprintf("Failed to parse conf: %s", je))
		}
		c.data = data
	}
	return c
}
<commit_msg>No nil map<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"sync"
)

type Conf struct {
	data    map[string]string
	dataMux sync.RWMutex
}

func (c *Conf) Get(k string) string {
	// Locking in base function
	return c.GetOrDefault(k, "")
}

func (c *Conf) Set(k string, v string) {
	c.dataMux.Lock()
	defer c.dataMux.Unlock()
	c.data[k] = v
}

func (c *Conf) Save() bool {
	c.dataMux.RLock()
	defer c.dataMux.RUnlock()
	b, je := json.Marshal(c.data)
	if je != nil {
		log.Printf("Failed saving conf: %s", je)
		return false
	}
	we := ioutil.WriteFile(confPath, b, 0600)
	if we != nil {
		log.Printf("Failed saving conf: %s", we)
		return false
	}
	return true
}

func (c *Conf) GetNotEmpty(k string) string {
	// Locking in base function
	val := c.GetOrDefault(k, "")
	if len(val) < 1 {
		panic(fmt.Sprintf("Value %s empty", k))
	}
	return val
}

func (c *Conf) GetOrDefault(k string, d string) string {
	c.dataMux.RLock()
	defer c.dataMux.RUnlock()
	if len(c.data[k]) == 0 {
		return d
	}
	return c.data[k]
}

func newConf(path string) *Conf {
	c := &Conf{}
	if len(path) > 0 {
		// Load file
		b, e := ioutil.ReadFile(path)
		if e != nil {
			log.Fatal(fmt.Sprintf("Failed to load conf: %s", e))
		}

		// Parse JSON
		var data map[string]string
		je := json.Unmarshal(b, &data)
		if je != nil {
			log.Fatal(fmt.Sprintf("Failed to parse conf: %s", je))
		}
		c.data = data
	} else {
		c.data = make(map[string]string)
	}
	return c
}
<|endoftext|>


<commit_before>package supervisor

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"golang.org/x/net/context"
)

func ExampleSupervisor() {
	var supervisor Supervisor

	svc := Simpleservice(1)
	supervisor.Add(&svc)

	ctx, _ := context.WithTimeout(context.Background(), 1*time.Second)
	supervisor.Serve(ctx)
}

func TestString(t *testing.T) {
	t.Parallel()

	const expected = "test"
	var supervisor Supervisor
	supervisor.Name = expected

	if got := fmt.Sprintf("%s", &supervisor); got != expected {
		t.Errorf("error getting supervisor name: %s", got)
	}
}

func TestCascaded(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := waitservice{id: 1}
	supervisor.Add(&svc1)
	svc2 := waitservice{id: 2}
	supervisor.Add(&svc2)

	var childSupervisor Supervisor
	svc3 := waitservice{id: 3}
	childSupervisor.Add(&svc3)
	svc4 := waitservice{id: 4}
	childSupervisor.Add(&svc4)

	supervisor.Add(&childSupervisor)

	ctx, _ := context.WithTimeout(context.Background(), 1*time.Second)
	supervisor.Serve(ctx)

	if count := getServiceCount(&supervisor); count != 3 {
		t.Errorf("unexpected service count: %v", count)
	}

	switch {
	case svc1.count != 1, svc2.count != 1, svc3.count != 1, svc4.count != 1:
		t.Errorf("services should have been executed only once. %d %d %d %d",
			svc1.count, svc2.count, svc3.count, svc4.count)
	}
}

func TestCascadedWithProblems(t *testing.T) {
	t.Parallel()

	supervisor := Supervisor{
		Backoff: 1 * time.Second,
		Log: func(msg string) {
			t.Log("supervisor log (cascaded with problems):", msg)
		},
	}
	svc1 := waitservice{id: 1}
	supervisor.Add(&svc1)
	svc2 := panicservice{id: 2}
	supervisor.Add(&svc2)

	childSupervisor := Supervisor{
		Backoff: 1 * time.Second,
	}
	svc3 := waitservice{id: 3}
	childSupervisor.Add(&svc3)
	svc4 := failingservice{id: 4}
	childSupervisor.Add(&svc4)

	supervisor.Add(&childSupervisor)

	ctx, _ := context.WithTimeout(context.Background(), 5*time.Second)
	supervisor.Serve(ctx)

	if count := getServiceCount(&supervisor); count != 3 {
		t.Errorf("unexpected service count: %v", count)
	}

	switch {
	case svc1.count != 1, svc3.count != 1:
		t.Errorf("services should have been executed only once. %d %d %d %d",
			svc1.count, svc2.count, svc3.count, svc4.count)
	case svc2.count <= 1, svc4.count <= 1:
		t.Errorf("services should have been executed at least once. %d %d %d %d",
			svc1.count, svc2.count, svc3.count, svc4.count)
	}
}

func TestPanic(t *testing.T) {
	t.Parallel()

	supervisor := Supervisor{
		Backoff: 500 * time.Millisecond,
		Log: func(msg string) {
			t.Log("supervisor log (panic):", msg)
		},
	}
	svc1 := panicservice{id: 1}
	supervisor.Add(&svc1)

	ctx, _ := context.WithTimeout(context.Background(), 1*time.Second)
	supervisor.Serve(ctx)

	// should arrive here with no panic
	if svc1.count == 1 {
		t.Error("the failed service should have been started at least once.")
	}
}

func TestFailing(t *testing.T) {
	t.Parallel()

	supervisor := Supervisor{
		Backoff: 1 * time.Second,
		Log: func(msg string) {
			t.Log("supervisor log (failing):", msg)
		},
	}

	svc1 := failingservice{id: 1}
	supervisor.Add(&svc1)

	ctx, _ := context.WithTimeout(context.Background(), 3*time.Second)
	supervisor.Serve(ctx)

	// should arrive here with no panic
	if svc1.count == 1 {
		t.Error("the failed service should have been started at least once.")
	}
}

func TestAddServiceAfterServe(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
		wg.Done()
	}()

	<-supervisor.startedServices
	svc2 := Simpleservice(2)
	supervisor.Add(&svc2)
	<-supervisor.startedServices

	cancel()
	<-ctx.Done()
	wg.Wait()

	if count := getServiceCount(&supervisor); count != 2 {
		t.Errorf("unexpected service count: %v", count)
	}
}

func TestRemoveServiceAfterServe(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)
	svc2 := Simpleservice(2)
	supervisor.Add(&svc2)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
		wg.Done()
	}()

	lbefore := getServiceCount(&supervisor)
	supervisor.Remove("unknown service")
	lafter := getServiceCount(&supervisor)

	if lbefore != lafter {
		t.Error("the removal of an unknown service shouldn't happen")
	}

	<-supervisor.startedServices
	supervisor.Remove(svc1.String())

	lremoved := getServiceCount(&supervisor)
	if lbefore != lremoved {
		t.Error("the removal of a service should have affected the supervisor:", lbefore, lremoved)
	}

	cancel()
	<-ctx.Done()
	wg.Wait()
}

func TestServices(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)
	svc2 := Simpleservice(2)
	supervisor.Add(&svc2)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
	}()

	<-supervisor.startedServices
	svcs := supervisor.Services()
	for _, svcname := range []string{svc1.String(), svc2.String()} {
		if _, ok := svcs[svcname]; !ok {
			t.Errorf("expected service not found: %s", svcname)
		}
	}

	cancel()
	<-ctx.Done()
	wg.Done()
}

func TestManualCancelation(t *testing.T) {
	t.Parallel()

	supervisor := Supervisor{
		Log: func(msg string) {
			t.Log("supervisor log (restartable):", msg)
		},
	}

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)
	svc2 := restartableservice{2, make(chan struct{})}
	supervisor.Add(&svc2)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
	}()

	<-supervisor.startedServices
	<-svc2.restarted

	// Testing restart
	svcs := supervisor.Cancelations()
	svcancel := svcs[svc2.String()]
	svcancel()
	<-svc2.restarted

	cancel()
	<-ctx.Done()
	wg.Done()

	// should arrive here with no panic
}

func TestServiceList(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
	}()

	<-supervisor.startedServices

	svcs := supervisor.Services()
	if svc, ok := svcs[svc1.String()]; !ok || &svc1 != svc.(*Simpleservice) {
		t.Errorf("could not find service when listing them. %s missing", svc1.String())
	}

	cancel()
	<-ctx.Done()
	wg.Done()
}

func TestDoubleStart(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	var svc1 waitservice
	supervisor.Add(&svc1)

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

	var wg sync.WaitGroup

	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(i int) {
			c := context.WithValue(ctx, "supervisor", i)
			supervisor.Serve(c)

			svc1.mu.Lock()
			count := svc1.count
			supervisors := svc1.supervisors
			if count > 1 {
				t.Error("wait service should have been started once:", count, "supervisor IDs:", supervisors)
			}
			svc1.mu.Unlock()

			wg.Done()
		}(i)
	}
	<-supervisor.startedServices

	cancel()
	<-ctx.Done()
	wg.Wait()
}

func TestRestart(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	var svc1 waitservice
	supervisor.Add(&svc1)

	for i := 1; i <= 3; i++ {
		ctx, _ := context.WithTimeout(context.Background(), 1*time.Second)
		supervisor.Serve(ctx)
		if svc1.count != i {
			t.Errorf("wait service should have been started %d. got: %d", i, svc1.count)
		}
	}
}

type failingservice struct {
	id, count int
}

func (s *failingservice) Serve(ctx context.Context) {
	select {
	case <-ctx.Done():
		return
	default:
		time.Sleep(100 * time.Millisecond)
		s.count++
		return
	}
}

func (s *failingservice) String() string {
	return fmt.Sprintf("failing service %v", s.id)
}

type panicservice struct {
	id, count int
}

func (s *panicservice) Serve(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
			time.Sleep(100 * time.Millisecond)
			s.count++
			panic("forcing panic")
		}
	}
}

func (s *panicservice) String() string {
	return fmt.Sprintf("panic service %v", s.id)
}

type restartableservice struct {
	id        int
	restarted chan struct{}
}

func (s *restartableservice) Serve(ctx context.Context) {
	var i int
	for {
		i++
		select {
		case <-ctx.Done():
			return
		default:
			time.Sleep(500 * time.Millisecond)
			select {
			case s.restarted <- struct{}{}:
			default:
			}
		}
	}
}

func (s *restartableservice) String() string {
	return fmt.Sprintf("restartable service %v", *s)
}

type Simpleservice int

func (s *Simpleservice) String() string {
	return fmt.Sprintf("simple service %d", int(*s))
}

func (s *Simpleservice) Serve(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
			time.Sleep(500 * time.Millisecond)
		}
	}
}

type waitservice struct {
	id          int
	mu          sync.Mutex
	count       int
	supervisors []int
}

func (s *waitservice) Serve(ctx context.Context) {
	s.mu.Lock()
	s.count++
	id := ctx.Value("supervisor")
	if id != nil {
		s.supervisors = append(s.supervisors, id.(int))
	}
	s.mu.Unlock()
	<-ctx.Done()
}

func (s *waitservice) String() string {
	return fmt.Sprintf("wait service %v", s.id)
}

func getServiceCount(s *Supervisor) int {
	s.servicesMu.Lock()
	l := len(s.services)
	s.servicesMu.Unlock()
	return l
}
<commit_msg>Add test for supervisor.Log<commit_after>package supervisor

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"golang.org/x/net/context"
)

func ExampleSupervisor() {
	var supervisor Supervisor

	svc := Simpleservice(1)
	supervisor.Add(&svc)

	ctx, _ := context.WithTimeout(context.Background(), 1*time.Second)
	supervisor.Serve(ctx)
}

func TestString(t *testing.T) {
	t.Parallel()

	const expected = "test"
	var supervisor Supervisor
	supervisor.Name = expected

	if got := fmt.Sprintf("%s", &supervisor); got != expected {
		t.Errorf("error getting supervisor name: %s", got)
	}
}

func TestCascaded(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := waitservice{id: 1}
	supervisor.Add(&svc1)
	svc2 := waitservice{id: 2}
	supervisor.Add(&svc2)

	var childSupervisor Supervisor
	svc3 := waitservice{id: 3}
	childSupervisor.Add(&svc3)
	svc4 := waitservice{id: 4}
	childSupervisor.Add(&svc4)

	supervisor.Add(&childSupervisor)

	ctx, _ := context.WithTimeout(context.Background(), 1*time.Second)
	supervisor.Serve(ctx)

	if count := getServiceCount(&supervisor); count != 3 {
		t.Errorf("unexpected service count: %v", count)
	}

	switch {
	case svc1.count != 1, svc2.count != 1, svc3.count != 1, svc4.count != 1:
		t.Errorf("services should have been executed only once. %d %d %d %d",
			svc1.count, svc2.count, svc3.count, svc4.count)
	}
}

func TestLog(t *testing.T) {
	t.Parallel()

	supervisor := Supervisor{
		Backoff: 1 * time.Second,
	}

	svc1 := panicservice{id: 1}
	supervisor.Add(&svc1)

	ctx, _ := context.WithTimeout(context.Background(), 5*time.Second)
	supervisor.Serve(ctx)
}

func TestCascadedWithProblems(t *testing.T) {
	t.Parallel()

	supervisor := Supervisor{
		Backoff: 1 * time.Second,
		Log: func(msg string) {
			t.Log("supervisor log (cascaded with problems):", msg)
		},
	}
	svc1 := waitservice{id: 1}
	supervisor.Add(&svc1)
	svc2 := panicservice{id: 2}
	supervisor.Add(&svc2)

	childSupervisor := Supervisor{
		Backoff: 1 * time.Second,
		Log: func(msg string) {
			t.Log("supervisor log (cascaded with problems - child):", msg)
		},
	}
	svc3 := waitservice{id: 3}
	childSupervisor.Add(&svc3)
	svc4 := failingservice{id: 4}
	childSupervisor.Add(&svc4)

	supervisor.Add(&childSupervisor)

	ctx, _ := context.WithTimeout(context.Background(), 5*time.Second)
	supervisor.Serve(ctx)

	if count := getServiceCount(&supervisor); count != 3 {
		t.Errorf("unexpected service count: %v", count)
	}

	switch {
	case svc1.count != 1, svc3.count != 1:
		t.Errorf("services should have been executed only once. %d %d %d %d",
			svc1.count, svc2.count, svc3.count, svc4.count)
	case svc2.count <= 1, svc4.count <= 1:
		t.Errorf("services should have been executed at least once. %d %d %d %d",
			svc1.count, svc2.count, svc3.count, svc4.count)
	}
}

func TestPanic(t *testing.T) {
	t.Parallel()

	supervisor := Supervisor{
		Backoff: 500 * time.Millisecond,
		Log: func(msg string) {
			t.Log("supervisor log (panic):", msg)
		},
	}
	svc1 := panicservice{id: 1}
	supervisor.Add(&svc1)

	ctx, _ := context.WithTimeout(context.Background(), 1*time.Second)
	supervisor.Serve(ctx)

	// should arrive here with no panic
	if svc1.count == 1 {
		t.Error("the failed service should have been started at least once.")
	}
}

func TestFailing(t *testing.T) {
	t.Parallel()

	supervisor := Supervisor{
		Backoff: 1 * time.Second,
		Log: func(msg string) {
			t.Log("supervisor log (failing):", msg)
		},
	}

	svc1 := failingservice{id: 1}
	supervisor.Add(&svc1)

	ctx, _ := context.WithTimeout(context.Background(), 3*time.Second)
	supervisor.Serve(ctx)

	// should arrive here with no panic
	if svc1.count == 1 {
		t.Error("the failed service should have been started at least once.")
	}
}

func TestAddServiceAfterServe(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
		wg.Done()
	}()

	<-supervisor.startedServices
	svc2 := Simpleservice(2)
	supervisor.Add(&svc2)
	<-supervisor.startedServices

	cancel()
	<-ctx.Done()
	wg.Wait()

	if count := getServiceCount(&supervisor); count != 2 {
		t.Errorf("unexpected service count: %v", count)
	}
}

func TestRemoveServiceAfterServe(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)
	svc2 := Simpleservice(2)
	supervisor.Add(&svc2)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
		wg.Done()
	}()

	lbefore := getServiceCount(&supervisor)
	supervisor.Remove("unknown service")
	lafter := getServiceCount(&supervisor)

	if lbefore != lafter {
		t.Error("the removal of an unknown service shouldn't happen")
	}

	<-supervisor.startedServices
	supervisor.Remove(svc1.String())

	lremoved := getServiceCount(&supervisor)
	if lbefore != lremoved {
		t.Error("the removal of a service should have affected the supervisor:", lbefore, lremoved)
	}

	cancel()
	<-ctx.Done()
	wg.Wait()
}

func TestServices(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)
	svc2 := Simpleservice(2)
	supervisor.Add(&svc2)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
	}()

	<-supervisor.startedServices
	svcs := supervisor.Services()
	for _, svcname := range []string{svc1.String(), svc2.String()} {
		if _, ok := svcs[svcname]; !ok {
			t.Errorf("expected service not found: %s", svcname)
		}
	}

	cancel()
	<-ctx.Done()
	wg.Done()
}

func TestManualCancelation(t *testing.T) {
	t.Parallel()

	supervisor := Supervisor{
		Log: func(msg string) {
			t.Log("supervisor log (restartable):", msg)
		},
	}

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)
	svc2 := restartableservice{2, make(chan struct{})}
	supervisor.Add(&svc2)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
	}()

	<-supervisor.startedServices
	<-svc2.restarted

	// Testing restart
	svcs := supervisor.Cancelations()
	svcancel := svcs[svc2.String()]
	svcancel()
	<-svc2.restarted

	cancel()
	<-ctx.Done()
	wg.Done()

	// should arrive here with no panic
}

func TestServiceList(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	svc1 := Simpleservice(1)
	supervisor.Add(&svc1)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		supervisor.Serve(ctx)
	}()

	<-supervisor.startedServices

	svcs := supervisor.Services()
	if svc, ok := svcs[svc1.String()]; !ok || &svc1 != svc.(*Simpleservice) {
		t.Errorf("could not find service when listing them. %s missing", svc1.String())
	}

	cancel()
	<-ctx.Done()
	wg.Done()
}

func TestDoubleStart(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	var svc1 waitservice
	supervisor.Add(&svc1)

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

	var wg sync.WaitGroup

	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(i int) {
			c := context.WithValue(ctx, "supervisor", i)
			supervisor.Serve(c)

			svc1.mu.Lock()
			count := svc1.count
			supervisors := svc1.supervisors
			if count > 1 {
				t.Error("wait service should have been started once:", count, "supervisor IDs:", supervisors)
			}
			svc1.mu.Unlock()

			wg.Done()
		}(i)
	}
	<-supervisor.startedServices

	cancel()
	<-ctx.Done()
	wg.Wait()
}

func TestRestart(t *testing.T) {
	t.Parallel()

	var supervisor Supervisor

	var svc1 waitservice
	supervisor.Add(&svc1)

	for i := 1; i <= 3; i++ {
		ctx, _ := context.WithTimeout(context.Background(), 1*time.Second)
		supervisor.Serve(ctx)
		if svc1.count != i {
			t.Errorf("wait service should have been started %d. got: %d", i, svc1.count)
		}
	}
}

type failingservice struct {
	id, count int
}

func (s *failingservice) Serve(ctx context.Context) {
	select {
	case <-ctx.Done():
		return
	default:
		time.Sleep(100 * time.Millisecond)
		s.count++
		return
	}
}

func (s *failingservice) String() string {
	return fmt.Sprintf("failing service %v", s.id)
}

type panicservice struct {
	id, count int
}

func (s *panicservice) Serve(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
			time.Sleep(100 * time.Millisecond)
			s.count++
			panic("forcing panic")
		}
	}
}

func (s *panicservice) String() string {
	return fmt.Sprintf("panic service %v", s.id)
}

type restartableservice struct {
	id        int
	restarted chan struct{}
}

func (s *restartableservice) Serve(ctx context.Context) {
	var i int
	for {
		i++
		select {
		case <-ctx.Done():
			return
		default:
			time.Sleep(500 * time.Millisecond)
			select {
			case s.restarted <- struct{}{}:
			default:
			}
		}
	}
}

func (s *restartableservice) String() string {
	return fmt.Sprintf("restartable service %v", *s)
}

type Simpleservice int

func (s *Simpleservice) String() string {
	return fmt.Sprintf("simple service %d", int(*s))
}

func (s *Simpleservice) Serve(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
			time.Sleep(500 * time.Millisecond)
		}
	}
}

type waitservice struct {
	id          int
	mu          sync.Mutex
	count       int
	supervisors []int
}

func (s *waitservice) Serve(ctx context.Context) {
	s.mu.Lock()
	s.count++
	id := ctx.Value("supervisor")
	if id != nil {
		s.supervisors = append(s.supervisors, id.(int))
	}
	s.mu.Unlock()
	<-ctx.Done()
}

func (s *waitservice) String() string {
	return fmt.Sprintf("wait service %v", s.id)
}

func getServiceCount(s *Supervisor) int {
	s.servicesMu.Lock()
	l := len(s.services)
	s.servicesMu.Unlock()
	return l
}
<|endoftext|>


<commit_before>package twitter

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/dghubble/sling"
)

const (
	userAgent    = "go-twitter v0.1"
	publicStream = "https://stream.twitter.com/1.1/"
	userStream   = "https://userstream.twitter.com/1.1/"
	siteStream   = "https://sitestream.twitter.com/1.1/"
)

// StreamService provides methods for accessing the Twitter Streaming API.
type StreamService struct {
	client *http.Client
	public *sling.Sling
	user   *sling.Sling
	site   *sling.Sling
}

// newStreamService returns a new StreamService.
func newStreamService(client *http.Client, sling *sling.Sling) *StreamService {
	sling.Set("User-Agent", userAgent)
	return &StreamService{
		client: client,
		public: sling.New().Base(publicStream).Path("statuses/"),
		user:   sling.New().Base(userStream),
		site:   sling.New().Base(siteStream),
	}
}

// StreamFilterParams are parameters for StreamService.Filter.
type StreamFilterParams struct {
	FilterLevel   string   `url:"filter_level,omitempty"`
	Follow        []string `url:"follow,omitempty,comma"`
	Language      []string `url:"language,omitempty,comma"`
	Locations     []string `url:"locations,omitempty,comma"`
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
	Track         []string `url:"track,omitempty,comma"`
}

// Filter returns messages that match one or more filter predicates.
// https://dev.twitter.com/streaming/reference/post/statuses/filter
func (srv *StreamService) Filter(params *StreamFilterParams) (*Stream, error) {
	req, err := srv.public.New().Post("filter.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// StreamSampleParams are the parameters for StreamService.Sample.
type StreamSampleParams struct {
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
	Language      []string `url:"language,omitempty,comma"`
}

// Sample returns a small sample of public stream messages.
// https://dev.twitter.com/streaming/reference/get/statuses/sample
func (srv *StreamService) Sample(params *StreamSampleParams) (*Stream, error) {
	req, err := srv.public.New().Get("sample.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// StreamUserParams are the parameters for StreamService.User.
type StreamUserParams struct {
	FilterLevel   string   `url:"filter_level,omitempty"`
	Language      []string `url:"language,omitempty,comma"`
	Locations     []string `url:"locations,omitempty,comma"`
	Replies       string   `url:"replies,omitempty"`
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
	Track         []string `url:"track,omitempty,comma"`
	With          string   `url:"with,omitempty"`
}

// User returns a stream of messages specific to the authenticated User.
// https://dev.twitter.com/streaming/reference/get/user
func (srv *StreamService) User(params *StreamUserParams) (*Stream, error) {
	req, err := srv.user.New().Get("user.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// StreamSiteParams are the parameters for StreamService.Site.
type StreamSiteParams struct {
	FilterLevel   string   `url:"filter_level,omitempty"`
	Follow        []string `url:"follow,omitempty,comma"`
	Language      []string `url:"language,omitempty,comma"`
	Replies       string   `url:"replies,omitempty"`
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
	With          string   `url:"with,omitempty"`
}

// Site returns messages for a set of users.
// Requires special permission to access.
// https://dev.twitter.com/streaming/reference/get/site
func (srv *StreamService) Site(params *StreamSiteParams) (*Stream, error) {
	req, err := srv.site.New().Get("site.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// StreamFirehoseParams are the parameters for StreamService.Firehose.
type StreamFirehoseParams struct {
	Count         int      `url:"count,omitempty"`
	FilterLevel   string   `url:"filter_level,omitempty"`
	Language      []string `url:"language,omitempty,comma"`
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
}

// Firehose returns all public messages and statuses.
// Requires special permission to access.
// https://dev.twitter.com/streaming/reference/get/statuses/firehose
func (srv *StreamService) Firehose(params *StreamFirehoseParams) (*Stream, error) {
	req, err := srv.public.New().Get("firehose.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// Stream maintains a connection to the Twitter Streaming API, receives
// messages from the streaming response, and sends them on the Messages
// channel from a goroutine. The stream goroutine stops itself if an EOF is
// reached or retry errors occur, also closing the Messages channel.
//
// The client must Stop() the stream when finished receiving, which will
// wait until the stream is properly stopped.
type Stream struct {
	client   *http.Client
	Messages chan interface{}
	done     chan struct{}
	group    *sync.WaitGroup
	body     io.Closer
}

// newStream creates a Stream and starts a goroutine to retry connecting and
// receive from a stream response. The goroutine may stop due to retry errors
// or be stopped by calling Stop() on the stream.
func newStream(client *http.Client, req *http.Request) *Stream {
	s := &Stream{
		client:   client,
		Messages: make(chan interface{}),
		done:     make(chan struct{}),
		group:    &sync.WaitGroup{},
	}
	s.group.Add(1)
	go s.retry(req, newExponentialBackOff(), newAggressiveExponentialBackOff())
	return s
}

// Stop signals retry and receiver to stop, closes the Messages channel, and
// blocks until done.
func (s *Stream) Stop() {
	close(s.done)
	// Scanner does not have a Stop() or take a done channel, so for low volume
	// streams Scan() blocks until the next keep-alive. Close the resp.Body to
	// escape and stop the stream in a timely fashion.
	if s.body != nil {
		s.body.Close()
	}
	// block until the retry goroutine stops
	s.group.Wait()
}

// retry retries making the given http.Request and receiving the response
// according to the Twitter backoff policies. Callers should invoke in a
// goroutine since backoffs sleep between retries.
// https://dev.twitter.com/streaming/overview/connecting
func (s *Stream) retry(req *http.Request, expBackOff backoff.BackOff, aggExpBackOff backoff.BackOff) {
	// close Messages channel and decrement the wait group counter
	defer close(s.Messages)
	defer s.group.Done()

	var wait time.Duration
	for !stopped(s.done) {
		resp, err := s.client.Do(req)
		if err != nil {
			// stop retrying for HTTP protocol errors
			s.Messages <- err
			return
		}
		// when err is nil, resp contains a non-nil Body which must be closed
		defer resp.Body.Close()
		s.body = resp.Body
		switch resp.StatusCode {
		case 200:
			// receive stream response Body, handles closing
			s.receive(resp.Body)
			expBackOff.Reset()
			aggExpBackOff.Reset()
		case 503:
			// exponential backoff
			wait = expBackOff.NextBackOff()
		case 420, 429:
			// aggressive exponential backoff
			wait = aggExpBackOff.NextBackOff()
		default:
			// stop retrying for other response codes
			resp.Body.Close()
			return
		}
		// close response before each retry
		resp.Body.Close()
		if wait == backoff.Stop {
			return
		}
		sleepOrDone(wait, s.done)
	}
}

// receive scans a stream response body, JSON decodes tokens to messages, and
// sends messages to the Messages channel. Receiving continues until an EOF,
// scan error, or the done channel is closed.
func (s *Stream) receive(body io.ReadCloser) {
	defer body.Close()
	reader := bufio.NewReader(body)
	for !stopped(s.done) {
		var buf []byte
		for {
			// Twitter streaming messages are separated with "\r\n", and a valid
			// message may sometimes contain '\n' in the middle.
			// bufio.Reader.Read() can accept one byte delimiter only, so we need to
			// first break out each line on '\n' and then check whether the line ends
			// with "\r\n" to find message boundaries.
			// https://dev.twitter.com/streaming/overview/processing
			line, err := reader.ReadBytes('\n')
			if err != nil {
				return
			}
			// If the line ends with "\r\n", it's the end of one streaming message data.
			if bytes.HasSuffix(line, []byte("\r\n")) {
				// reader.ReadBytes() returns a slice including the delimiter itself, so we
				// need to trim '\n' as well as '\r' from the end of the slice.
				buf = append(buf, bytes.TrimRight(line, "\r\n")...)
				break
			}
			// Otherwise, the line is not the end of a streaming message, so we append
			// the line to the buffer and continue to scan lines.
			buf = append(buf, line...)
		}
		if len(buf) == 0 {
			// empty keep-alive
			continue
		}
		select {
		// send messages, data, or errors
		case s.Messages <- getMessage(buf):
			continue
		// allow client to Stop(), even if not receiving
		case <-s.done:
			return
		}
	}
}

// getMessage unmarshals the token and returns a message struct, if the type
// can be determined. Otherwise, returns the token unmarshalled into a data
// map[string]interface{} or the unmarshal error.
func getMessage(token []byte) interface{} {
	var data map[string]interface{}
	// unmarshal JSON encoded token into a map for
	err := json.Unmarshal(token, &data)
	if err != nil {
		return err
	}
	return decodeMessage(token, data)
}

// decodeMessage determines the message type from known data keys, allocates
// at most one message struct, and JSON decodes the token into the message.
// Returns the message struct or the data map if the message type could not be
// determined.
func decodeMessage(token []byte, data map[string]interface{}) interface{} {
	if hasPath(data, "retweet_count") {
		tweet := new(Tweet)
		json.Unmarshal(token, tweet)
		return tweet
	} else if hasPath(data, "direct_message") {
		notice := new(directMessageNotice)
		json.Unmarshal(token, notice)
		return notice.DirectMessage
	} else if hasPath(data, "delete") {
		notice := new(statusDeletionNotice)
		json.Unmarshal(token, notice)
		return notice.Delete.StatusDeletion
	} else if hasPath(data, "scrub_geo") {
		notice := new(locationDeletionNotice)
		json.Unmarshal(token, notice)
		return notice.ScrubGeo
	} else if hasPath(data, "limit") {
		notice := new(streamLimitNotice)
		json.Unmarshal(token, notice)
		return notice.Limit
	} else if hasPath(data, "status_withheld") {
		notice := new(statusWithheldNotice)
		json.Unmarshal(token, notice)
		return notice.StatusWithheld
	} else if hasPath(data, "user_withheld") {
		notice := new(userWithheldNotice)
		json.Unmarshal(token, notice)
		return notice.UserWithheld
	} else if hasPath(data, "disconnect") {
		notice := new(streamDisconnectNotice)
		json.Unmarshal(token, notice)
		return notice.StreamDisconnect
	} else if hasPath(data, "warning") {
		notice := new(stallWarningNotice)
		json.Unmarshal(token, notice)
		return notice.StallWarning
	} else if hasPath(data, "friends") {
		friendsList := new(FriendsList)
		json.Unmarshal(token, friendsList)
		return friendsList
	} else if hasPath(data, "event") {
		event := new(Event)
		json.Unmarshal(token, event)
		return event
	}
	// message type unknown, return the data map[string]interface{}
	return data
}

// hasPath returns true if the map contains the given key, false otherwise.
func hasPath(data map[string]interface{}, key string) bool {
	_, ok := data[key]
	return ok
}
<commit_msg>Use bytes.Buffer for reading streaming messages<commit_after>package twitter

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/dghubble/sling"
)

const (
	userAgent    = "go-twitter v0.1"
	publicStream = "https://stream.twitter.com/1.1/"
	userStream   = "https://userstream.twitter.com/1.1/"
	siteStream   = "https://sitestream.twitter.com/1.1/"
)

// StreamService provides methods for accessing the Twitter Streaming API.
type StreamService struct {
	client *http.Client
	public *sling.Sling
	user   *sling.Sling
	site   *sling.Sling
}

// newStreamService returns a new StreamService.
func newStreamService(client *http.Client, sling *sling.Sling) *StreamService {
	sling.Set("User-Agent", userAgent)
	return &StreamService{
		client: client,
		public: sling.New().Base(publicStream).Path("statuses/"),
		user:   sling.New().Base(userStream),
		site:   sling.New().Base(siteStream),
	}
}

// StreamFilterParams are parameters for StreamService.Filter.
type StreamFilterParams struct {
	FilterLevel   string   `url:"filter_level,omitempty"`
	Follow        []string `url:"follow,omitempty,comma"`
	Language      []string `url:"language,omitempty,comma"`
	Locations     []string `url:"locations,omitempty,comma"`
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
	Track         []string `url:"track,omitempty,comma"`
}

// Filter returns messages that match one or more filter predicates.
// https://dev.twitter.com/streaming/reference/post/statuses/filter
func (srv *StreamService) Filter(params *StreamFilterParams) (*Stream, error) {
	req, err := srv.public.New().Post("filter.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// StreamSampleParams are the parameters for StreamService.Sample.
type StreamSampleParams struct {
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
	Language      []string `url:"language,omitempty,comma"`
}

// Sample returns a small sample of public stream messages.
// https://dev.twitter.com/streaming/reference/get/statuses/sample
func (srv *StreamService) Sample(params *StreamSampleParams) (*Stream, error) {
	req, err := srv.public.New().Get("sample.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// StreamUserParams are the parameters for StreamService.User.
type StreamUserParams struct {
	FilterLevel   string   `url:"filter_level,omitempty"`
	Language      []string `url:"language,omitempty,comma"`
	Locations     []string `url:"locations,omitempty,comma"`
	Replies       string   `url:"replies,omitempty"`
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
	Track         []string `url:"track,omitempty,comma"`
	With          string   `url:"with,omitempty"`
}

// User returns a stream of messages specific to the authenticated User.
// https://dev.twitter.com/streaming/reference/get/user
func (srv *StreamService) User(params *StreamUserParams) (*Stream, error) {
	req, err := srv.user.New().Get("user.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// StreamSiteParams are the parameters for StreamService.Site.
type StreamSiteParams struct {
	FilterLevel   string   `url:"filter_level,omitempty"`
	Follow        []string `url:"follow,omitempty,comma"`
	Language      []string `url:"language,omitempty,comma"`
	Replies       string   `url:"replies,omitempty"`
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
	With          string   `url:"with,omitempty"`
}

// Site returns messages for a set of users.
// Requires special permission to access.
// https://dev.twitter.com/streaming/reference/get/site
func (srv *StreamService) Site(params *StreamSiteParams) (*Stream, error) {
	req, err := srv.site.New().Get("site.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// StreamFirehoseParams are the parameters for StreamService.Firehose.
type StreamFirehoseParams struct {
	Count         int      `url:"count,omitempty"`
	FilterLevel   string   `url:"filter_level,omitempty"`
	Language      []string `url:"language,omitempty,comma"`
	StallWarnings *bool    `url:"stall_warnings,omitempty"`
}

// Firehose returns all public messages and statuses.
// Requires special permission to access.
// https://dev.twitter.com/streaming/reference/get/statuses/firehose
func (srv *StreamService) Firehose(params *StreamFirehoseParams) (*Stream, error) {
	req, err := srv.public.New().Get("firehose.json").QueryStruct(params).Request()
	if err != nil {
		return nil, err
	}
	return newStream(srv.client, req), nil
}

// Stream maintains a connection to the Twitter Streaming API, receives
// messages from the streaming response, and sends them on the Messages
// channel from a goroutine. The stream goroutine stops itself if an EOF is
// reached or retry errors occur, also closing the Messages channel.
//
// The client must Stop() the stream when finished receiving, which will
// wait until the stream is properly stopped.
type Stream struct {
	client   *http.Client
	Messages chan interface{}
	done     chan struct{}
	group    *sync.WaitGroup
	body     io.Closer
}

// newStream creates a Stream and starts a goroutine to retry connecting and
// receive from a stream response. The goroutine may stop due to retry errors
// or be stopped by calling Stop() on the stream.
func newStream(client *http.Client, req *http.Request) *Stream {
	s := &Stream{
		client:   client,
		Messages: make(chan interface{}),
		done:     make(chan struct{}),
		group:    &sync.WaitGroup{},
	}
	s.group.Add(1)
	go s.retry(req, newExponentialBackOff(), newAggressiveExponentialBackOff())
	return s
}

// Stop signals retry and receiver to stop, closes the Messages channel, and
// blocks until done.
func (s *Stream) Stop() {
	close(s.done)
	// Scanner does not have a Stop() or take a done channel, so for low volume
	// streams Scan() blocks until the next keep-alive. Close the resp.Body to
	// escape and stop the stream in a timely fashion.
	if s.body != nil {
		s.body.Close()
	}
	// block until the retry goroutine stops
	s.group.Wait()
}

// retry retries making the given http.Request and receiving the response
// according to the Twitter backoff policies. Callers should invoke in a
// goroutine since backoffs sleep between retries.
// https://dev.twitter.com/streaming/overview/connecting
func (s *Stream) retry(req *http.Request, expBackOff backoff.BackOff, aggExpBackOff backoff.BackOff) {
	// close Messages channel and decrement the wait group counter
	defer close(s.Messages)
	defer s.group.Done()

	var wait time.Duration
	for !stopped(s.done) {
		resp, err := s.client.Do(req)
		if err != nil {
			// stop retrying for HTTP protocol errors
			s.Messages <- err
			return
		}
		// when err is nil, resp contains a non-nil Body which must be closed
		defer resp.Body.Close()
		s.body = resp.Body
		switch resp.StatusCode {
		case 200:
			// receive stream response Body, handles closing
			s.receive(resp.Body)
			expBackOff.Reset()
			aggExpBackOff.Reset()
		case 503:
			// exponential backoff
			wait = expBackOff.NextBackOff()
		case 420, 429:
			// aggressive exponential backoff
			wait = aggExpBackOff.NextBackOff()
		default:
			// stop retrying for other response codes
			resp.Body.Close()
			return
		}
		// close response before each retry
		resp.Body.Close()
		if wait == backoff.Stop {
			return
		}
		sleepOrDone(wait, s.done)
	}
}

// receive scans a stream response body, JSON decodes tokens to messages, and
// sends messages to the Messages channel. Receiving continues until an EOF,
// scan error, or the done channel is closed.
func (s *Stream) receive(body io.ReadCloser) {
	defer body.Close()
	reader := bufio.NewReader(body)
	var buf bytes.Buffer
	for !stopped(s.done) {
		// Discard all the bytes from buf and continue to use the allocated memory
		// space for reading the next message.
		buf.Truncate(0)
		for {
			// Twitter streaming messages are separated with "\r\n", and a valid
			// message may sometimes contain '\n' in the middle.
			// bufio.Reader.Read() can accept one byte delimiter only, so we need to
			// first break out each line on '\n' and then check whether the line ends
			// with "\r\n" to find message boundaries.
			// https://dev.twitter.com/streaming/overview/processing
			line, err := reader.ReadBytes('\n')
			if err != nil {
				return
			}
			// If the line ends with "\r\n", it's the end of one streaming message data.
			if bytes.HasSuffix(line, []byte("\r\n")) {
				// reader.ReadBytes() returns a slice including the delimiter itself, so we
				// need to trim '\n' as well as '\r' from the end of the slice.
				buf.Write(bytes.TrimRight(line, "\r\n"))
				break
			}
			// Otherwise, the line is not the end of a streaming message, so we append
			// the line to the buffer and continue to scan lines.
			buf.Write(line)
		}
		// Get the streaming message bytes from buf. Not that Bytes() won't mark the
		// returned data as "read", and we need to explicitly call Truncate(0) to
		// discard from buf before writing the next streaming message to buf.
		data := buf.Bytes()
		if len(data) == 0 {
			// empty keep-alive
			continue
		}
		select {
		// send messages, data, or errors
		case s.Messages <- getMessage(data):
			continue
		// allow client to Stop(), even if not receiving
		case <-s.done:
			return
		}
	}
}

// getMessage unmarshals the token and returns a message struct, if the type
// can be determined. Otherwise, returns the token unmarshalled into a data
// map[string]interface{} or the unmarshal error.
func getMessage(token []byte) interface{} {
	var data map[string]interface{}
	// unmarshal JSON encoded token into a map for
	err := json.Unmarshal(token, &data)
	if err != nil {
		return err
	}
	return decodeMessage(token, data)
}

// decodeMessage determines the message type from known data keys, allocates
// at most one message struct, and JSON decodes the token into the message.
// Returns the message struct or the data map if the message type could not be
// determined.
func decodeMessage(token []byte, data map[string]interface{}) interface{} {
	if hasPath(data, "retweet_count") {
		tweet := new(Tweet)
		json.Unmarshal(token, tweet)
		return tweet
	} else if hasPath(data, "direct_message") {
		notice := new(directMessageNotice)
		json.Unmarshal(token, notice)
		return notice.DirectMessage
	} else if hasPath(data, "delete") {
		notice := new(statusDeletionNotice)
		json.Unmarshal(token, notice)
		return notice.Delete.StatusDeletion
	} else if hasPath(data, "scrub_geo") {
		notice := new(locationDeletionNotice)
		json.Unmarshal(token, notice)
		return notice.ScrubGeo
	} else if hasPath(data, "limit") {
		notice := new(streamLimitNotice)
		json.Unmarshal(token, notice)
		return notice.Limit
	} else if hasPath(data, "status_withheld") {
		notice := new(statusWithheldNotice)
		json.Unmarshal(token, notice)
		return notice.StatusWithheld
	} else if hasPath(data, "user_withheld") {
		notice := new(userWithheldNotice)
		json.Unmarshal(token, notice)
		return notice.UserWithheld
	} else if hasPath(data, "disconnect") {
		notice := new(streamDisconnectNotice)
		json.Unmarshal(token, notice)
		return notice.StreamDisconnect
	} else if hasPath(data, "warning") {
		notice := new(stallWarningNotice)
		json.Unmarshal(token, notice)
		return notice.StallWarning
	} else if hasPath(data, "friends") {
		friendsList := new(FriendsList)
		json.Unmarshal(token, friendsList)
		return friendsList
	} else if hasPath(data, "event") {
		event := new(Event)
		json.Unmarshal(token, event)
		return event
	}
	// message type unknown, return the data map[string]interface{}
	return data
}

// hasPath returns true if the map contains the given key, false otherwise.
func hasPath(data map[string]interface{}, key string) bool {
	_, ok := data[key]
	return ok
}
<|endoftext|>


<commit_before>package adodb

import (
	"fmt"
	"math/big"
	"errors"
	"exp/sql"
	"exp/sql/driver"
	"github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
	"time"
	"unsafe"
)

func init() {
	ole.CoInitialize(0)
	sql.Register("adodb", &AdodbDriver{})
}

type AdodbDriver struct {

}

type AdodbConn struct {
	db *ole.IDispatch
}

type AdodbTx struct {
	c *AdodbConn
}

func (tx *AdodbTx) Commit() error {
	_, err := oleutil.CallMethod(tx.c.db, "CommitTrans")
	if err != nil {
		return err
	}
	return nil
}

func (tx *AdodbTx) Rollback() error {
	_, err := oleutil.CallMethod(tx.c.db, "Rollback")
	if err != nil {
		return err
	}
	return nil
}

func (c *AdodbConn) exec(cmd string) error {
	_, err := oleutil.CallMethod(c.db, "Execute", cmd)
	return err
}

func (c *AdodbConn) Begin() (driver.Tx, error) {
	_, err := oleutil.CallMethod(c.db, "BeginTrans")
	if err != nil {
		return nil, err
	}
	return &AdodbTx{c}, nil
}

func (d *AdodbDriver) Open(dsn string) (driver.Conn, error) {
	unknown, err := oleutil.CreateObject("ADODB.Connection")
	if err != nil {
		return nil, err
	}
	db, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.CallMethod(db, "Open", dsn)
	if err != nil {
		return nil, err
	}
	return &AdodbConn{db}, nil
}

func (c *AdodbConn) Close() error {
	_, err := oleutil.CallMethod(c.db, "Close")
	if err != nil {
		return err
	}
	c.db = nil
	return nil
}

type AdodbStmt struct {
	c *AdodbConn
	s *ole.IDispatch
	ps *ole.IDispatch
	b []string
}

func (c *AdodbConn) Prepare(query string) (driver.Stmt, error) {
	unknown, err := oleutil.CreateObject("ADODB.Command")
	if err != nil {
		return nil, err
	}
	s, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "ActiveConnection", c.db)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandText", query)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandType", 1)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "Prepared", true)
	if err != nil {
		return nil, err
	}
	val, err := oleutil.GetProperty(s, "Parameters")
	if err != nil {
		return nil, err
	}
	return &AdodbStmt{c, s, val.ToIDispatch(), nil}, nil
}

func (s *AdodbStmt) Bind(bind []string) error {
	s.b = bind
	return nil
}

func (s *AdodbStmt) Close() error {
	s.s.Release()
	return nil
}

func (s *AdodbStmt) NumInput() int {
	if s.b != nil {
		return len(s.b)
	}
	_, err := oleutil.CallMethod(s.ps, "Refresh")
	if err != nil {
		return -1
	}
	val, err := oleutil.GetProperty(s.ps, "Count")
	if err != nil {
		return -1
	}
	c := int(val.Val)
	return c
}

func (s *AdodbStmt) bind(args []interface{}) error {
	if s.b != nil {
		for i, v := range args {
			var b string = "?"
			if len(s.b) < i {
				b = s.b[i]
			}
			unknown, err := oleutil.CallMethod(s.s, "CreateParameter", b, 12, 1)
			if err != nil {
				return err
			}
			param := unknown.ToIDispatch()
			defer param.Release()
			_, err = oleutil.PutProperty(param, "Value", v)
			if err != nil {
				return err
			}
			_, err = oleutil.CallMethod(s.ps, "Append", param)
			if err != nil {
				return err
			}
		}
	} else {
		for i, v := range args {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(s.ps, "Item", &varval)
			if err != nil {
				return err
			}
			item := val.ToIDispatch()
			defer item.Release()
			_, err = oleutil.PutProperty(item, "Value", v)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *AdodbStmt) Query(args []interface{}) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rc, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return &AdodbRows{s, rc.ToIDispatch(), -1, nil}, nil
}

func (s *AdodbStmt) Exec(args []interface{}) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	_, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return driver.DDLSuccess, nil
}

type AdodbRows struct {
	s    *AdodbStmt
	rc   *ole.IDispatch
	nc   int
	cols []string
}

func (rc *AdodbRows) Close() error {
	_, err := oleutil.CallMethod(rc.rc, "Close")
	if err != nil {
		return err
	}
	return nil
}

func (rc *AdodbRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		unknown, err := oleutil.GetProperty(rc.rc, "Fields")
		if err != nil {
			return []string {}
		}
		fields := unknown.ToIDispatch()
		defer fields.Release()
		val, err := oleutil.GetProperty(fields, "Count")
		if err != nil {
			return []string {}
		}
		rc.nc = int(val.Val)
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(fields, "Item", &varval)
			if err != nil {
				return []string {}
			}
			item := val.ToIDispatch()
			if err != nil {
				return []string {}
			}
			name, err := oleutil.GetProperty(item, "Name")
			if err != nil {
				return []string {}
			}
			rc.cols[i] = name.ToString()
			item.Release()
		}
	}
	return rc.cols
}

func (rc *AdodbRows) Next(dest []interface{}) error {
	_, err := oleutil.CallMethod(rc.rc, "MoveNext")
	if err != nil {
		fmt.Println(err)
		return err
	}
	unknown, err := oleutil.GetProperty(rc.rc, "EOF")
	if err != nil {
		return err
	}
	if unknown.Val != 0 {
		return errors.New("EOF")
	}
	unknown, err = oleutil.GetProperty(rc.rc, "Fields")
	if err != nil {
		return err
	}
	fields := unknown.ToIDispatch()
	defer fields.Release()
	for i := range dest {
		var varval ole.VARIANT
		varval.VT = ole.VT_I4
		varval.Val = int64(i)
		val, err := oleutil.CallMethod(fields, "Item", &varval)
		if err != nil {
			return err
		}
		field := val.ToIDispatch()
		defer field.Release()
		typ, err := oleutil.GetProperty(field, "Type")
		if err != nil {
			return err
		}
		val, err = oleutil.GetProperty(field, "Value")
		if err != nil {
			return err
		}
		field.Release()
		switch typ.Val {
		case 0: // ADEMPTY
			// TODO
		case 2: // ADSMALLINT
			dest[i] = int16(val.Val)
		case 3: // ADINTEGER
			dest[i] = int32(val.Val)
		case 4: // ADSINGLE
			dest[i] = float32(val.Val)
		case 5: // ADDOUBLE
			dest[i] = val.Val
		case 6: // ADCURRENCY
			// TODO
		case 7: // ADDATE
			// TODO
		case 8: // ADBSTR
			dest[i] = val.ToString()
		case 9: // ADIDISPATCH
			dest[i] = val.ToIDispatch()
		case 10: // ADERROR
			// TODO
		case 11: // ADBOOLEAN
			if val.Val != 0 {
				dest[i] = true
			} else {
				dest[i] = false
			}
		case 12: // ADVARIANT
			dest[i] = val
		case 13: // ADIUNKNOWN
			dest[i] = val.ToIUnknown()
		case 14: // ADDECIMAL
			// TODO
		case 16: // ADTINYINT
			dest[i] = int8(val.Val)
		case 17: // ADUNSIGNEDTINYINT
			dest[i] = uint8(val.Val)
		case 18: // ADUNSIGNEDSMALLINT
			dest[i] = uint16(val.Val)
		case 19: // ADUNSIGNEDINT
			dest[i] = uint32(val.Val)
		case 20: // ADBIGINT
			dest[i] = big.NewInt(val.Val)
		case 21: // ADUNSIGNEDBIGINT
			// TODO
		case 72: // ADGUID
			// TODO
		case 128: // ADBINARY
			sa := *(**ole.SAFEARRAY)(unsafe.Pointer(&val.Val))
			dest[i] = (*[1 << 30]byte)(unsafe.Pointer(uintptr(sa.PvData)))[0:sa.CbElements]
		case 129: // ADCHAR
			dest[i] = uint8(val.Val)
		case 130: // ADWCHAR
			dest[i] = uint16(val.Val)
		case 131: // ADNUMERIC
			dest[i] = val.Val
		case 132: // ADUSERDEFINED
			dest[i] = uintptr(val.Val)
		case 133: // ADDBDATE
			dest[i] = time. NanosecondsToUTC(val.Val)
		case 134: // ADDBTIME
			dest[i] = time. NanosecondsToUTC(val.Val)
		case 135: // ADDBTIMESTAMP
			dest[i] = time. NanosecondsToUTC(val.Val)
		case 136: // ADCHAPTER
			dest[i] = val.ToString()
		case 200: // ADVARCHAR
			dest[i] = val.ToString()
		case 201: // ADLONGVARCHAR
			dest[i] = val.ToString()
		case 202: // ADVARWCHAR
			dest[i] = val.ToString()
		case 203: // ADLONGVARWCHAR
			dest[i] = val.ToString()
		case 204: // ADVARBINARY
			// TODO
		case 205: // ADLONGVARBINARY
			// TODO
		}
	}
	return nil
}
<commit_msg>gofix.<commit_after>package adodb

import (
	"errors"
	"exp/sql"
	"exp/sql/driver"
	"fmt"
	"github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
	"math/big"
	"time"
	"unsafe"
)

func init() {
	ole.CoInitialize(0)
	sql.Register("adodb", &AdodbDriver{})
}

type AdodbDriver struct {

}

type AdodbConn struct {
	db *ole.IDispatch
}

type AdodbTx struct {
	c *AdodbConn
}

func (tx *AdodbTx) Commit() error {
	_, err := oleutil.CallMethod(tx.c.db, "CommitTrans")
	if err != nil {
		return err
	}
	return nil
}

func (tx *AdodbTx) Rollback() error {
	_, err := oleutil.CallMethod(tx.c.db, "Rollback")
	if err != nil {
		return err
	}
	return nil
}

func (c *AdodbConn) exec(cmd string) error {
	_, err := oleutil.CallMethod(c.db, "Execute", cmd)
	return err
}

func (c *AdodbConn) Begin() (driver.Tx, error) {
	_, err := oleutil.CallMethod(c.db, "BeginTrans")
	if err != nil {
		return nil, err
	}
	return &AdodbTx{c}, nil
}

func (d *AdodbDriver) Open(dsn string) (driver.Conn, error) {
	unknown, err := oleutil.CreateObject("ADODB.Connection")
	if err != nil {
		return nil, err
	}
	db, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.CallMethod(db, "Open", dsn)
	if err != nil {
		return nil, err
	}
	return &AdodbConn{db}, nil
}

func (c *AdodbConn) Close() error {
	_, err := oleutil.CallMethod(c.db, "Close")
	if err != nil {
		return err
	}
	c.db = nil
	return nil
}

type AdodbStmt struct {
	c  *AdodbConn
	s  *ole.IDispatch
	ps *ole.IDispatch
	b  []string
}

func (c *AdodbConn) Prepare(query string) (driver.Stmt, error) {
	unknown, err := oleutil.CreateObject("ADODB.Command")
	if err != nil {
		return nil, err
	}
	s, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "ActiveConnection", c.db)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandText", query)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandType", 1)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "Prepared", true)
	if err != nil {
		return nil, err
	}
	val, err := oleutil.GetProperty(s, "Parameters")
	if err != nil {
		return nil, err
	}
	return &AdodbStmt{c, s, val.ToIDispatch(), nil}, nil
}

func (s *AdodbStmt) Bind(bind []string) error {
	s.b = bind
	return nil
}

func (s *AdodbStmt) Close() error {
	s.s.Release()
	return nil
}

func (s *AdodbStmt) NumInput() int {
	if s.b != nil {
		return len(s.b)
	}
	_, err := oleutil.CallMethod(s.ps, "Refresh")
	if err != nil {
		return -1
	}
	val, err := oleutil.GetProperty(s.ps, "Count")
	if err != nil {
		return -1
	}
	c := int(val.Val)
	return c
}

func (s *AdodbStmt) bind(args []interface{}) error {
	if s.b != nil {
		for i, v := range args {
			var b string = "?"
			if len(s.b) < i {
				b = s.b[i]
			}
			unknown, err := oleutil.CallMethod(s.s, "CreateParameter", b, 12, 1)
			if err != nil {
				return err
			}
			param := unknown.ToIDispatch()
			defer param.Release()
			_, err = oleutil.PutProperty(param, "Value", v)
			if err != nil {
				return err
			}
			_, err = oleutil.CallMethod(s.ps, "Append", param)
			if err != nil {
				return err
			}
		}
	} else {
		for i, v := range args {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(s.ps, "Item", &varval)
			if err != nil {
				return err
			}
			item := val.ToIDispatch()
			defer item.Release()
			_, err = oleutil.PutProperty(item, "Value", v)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *AdodbStmt) Query(args []interface{}) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rc, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return &AdodbRows{s, rc.ToIDispatch(), -1, nil}, nil
}

func (s *AdodbStmt) Exec(args []interface{}) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	_, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return driver.DDLSuccess, nil
}

type AdodbRows struct {
	s    *AdodbStmt
	rc   *ole.IDispatch
	nc   int
	cols []string
}

func (rc *AdodbRows) Close() error {
	_, err := oleutil.CallMethod(rc.rc, "Close")
	if err != nil {
		return err
	}
	return nil
}

func (rc *AdodbRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		unknown, err := oleutil.GetProperty(rc.rc, "Fields")
		if err != nil {
			return []string{}
		}
		fields := unknown.ToIDispatch()
		defer fields.Release()
		val, err := oleutil.GetProperty(fields, "Count")
		if err != nil {
			return []string{}
		}
		rc.nc = int(val.Val)
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(fields, "Item", &varval)
			if err != nil {
				return []string{}
			}
			item := val.ToIDispatch()
			if err != nil {
				return []string{}
			}
			name, err := oleutil.GetProperty(item, "Name")
			if err != nil {
				return []string{}
			}
			rc.cols[i] = name.ToString()
			item.Release()
		}
	}
	return rc.cols
}

func (rc *AdodbRows) Next(dest []interface{}) error {
	_, err := oleutil.CallMethod(rc.rc, "MoveNext")
	if err != nil {
		fmt.Println(err)
		return err
	}
	unknown, err := oleutil.GetProperty(rc.rc, "EOF")
	if err != nil {
		return err
	}
	if unknown.Val != 0 {
		return errors.New("EOF")
	}
	unknown, err = oleutil.GetProperty(rc.rc, "Fields")
	if err != nil {
		return err
	}
	fields := unknown.ToIDispatch()
	defer fields.Release()
	for i := range dest {
		var varval ole.VARIANT
		varval.VT = ole.VT_I4
		varval.Val = int64(i)
		val, err := oleutil.CallMethod(fields, "Item", &varval)
		if err != nil {
			return err
		}
		field := val.ToIDispatch()
		defer field.Release()
		typ, err := oleutil.GetProperty(field, "Type")
		if err != nil {
			return err
		}
		val, err = oleutil.GetProperty(field, "Value")
		if err != nil {
			return err
		}
		field.Release()
		switch typ.Val {
		case 0: // ADEMPTY
			// TODO
		case 2: // ADSMALLINT
			dest[i] = int16(val.Val)
		case 3: // ADINTEGER
			dest[i] = int32(val.Val)
		case 4: // ADSINGLE
			dest[i] = float32(val.Val)
		case 5: // ADDOUBLE
			dest[i] = val.Val
		case 6: // ADCURRENCY
			// TODO
		case 7: // ADDATE
			// TODO
		case 8: // ADBSTR
			dest[i] = val.ToString()
		case 9: // ADIDISPATCH
			dest[i] = val.ToIDispatch()
		case 10: // ADERROR
			// TODO
		case 11: // ADBOOLEAN
			if val.Val != 0 {
				dest[i] = true
			} else {
				dest[i] = false
			}
		case 12: // ADVARIANT
			dest[i] = val
		case 13: // ADIUNKNOWN
			dest[i] = val.ToIUnknown()
		case 14: // ADDECIMAL
			// TODO
		case 16: // ADTINYINT
			dest[i] = int8(val.Val)
		case 17: // ADUNSIGNEDTINYINT
			dest[i] = uint8(val.Val)
		case 18: // ADUNSIGNEDSMALLINT
			dest[i] = uint16(val.Val)
		case 19: // ADUNSIGNEDINT
			dest[i] = uint32(val.Val)
		case 20: // ADBIGINT
			dest[i] = big.NewInt(val.Val)
		case 21: // ADUNSIGNEDBIGINT
			// TODO
		case 72: // ADGUID
			// TODO
		case 128: // ADBINARY
			sa := *(**ole.SAFEARRAY)(unsafe.Pointer(&val.Val))
			dest[i] = (*[1 << 30]byte)(unsafe.Pointer(uintptr(sa.PvData)))[0:sa.CbElements]
		case 129: // ADCHAR
			dest[i] = uint8(val.Val)
		case 130: // ADWCHAR
			dest[i] = uint16(val.Val)
		case 131: // ADNUMERIC
			dest[i] = val.Val
		case 132: // ADUSERDEFINED
			dest[i] = uintptr(val.Val)
		case 133: // ADDBDATE
			dest[i] = time.Unix(0, val.Val).UTC()
		case 134: // ADDBTIME
			dest[i] = time.Unix(0, val.Val).UTC()
		case 135: // ADDBTIMESTAMP
			dest[i] = time.Unix(0, val.Val).UTC()
		case 136: // ADCHAPTER
			dest[i] = val.ToString()
		case 200: // ADVARCHAR
			dest[i] = val.ToString()
		case 201: // ADLONGVARCHAR
			dest[i] = val.ToString()
		case 202: // ADVARWCHAR
			dest[i] = val.ToString()
		case 203: // ADLONGVARWCHAR
			dest[i] = val.ToString()
		case 204: // ADVARBINARY
			// TODO
		case 205: // ADLONGVARBINARY
			// TODO
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package adodb

import (
	"errors"
	"database/sql"
	"database/sql/driver"
	"github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
	"io"
	"math/big"
	"time"
	"unsafe"
)

func init() {
	ole.CoInitialize(0)
	sql.Register("adodb", &AdodbDriver{})
}

type AdodbDriver struct {

}

type AdodbConn struct {
	db *ole.IDispatch
}

type AdodbTx struct {
	c *AdodbConn
}

func (tx *AdodbTx) Commit() error {
	_, err := oleutil.CallMethod(tx.c.db, "CommitTrans")
	if err != nil {
		return err
	}
	return nil
}

func (tx *AdodbTx) Rollback() error {
	_, err := oleutil.CallMethod(tx.c.db, "Rollback")
	if err != nil {
		return err
	}
	return nil
}

func (c *AdodbConn) exec(cmd string) error {
	_, err := oleutil.CallMethod(c.db, "Execute", cmd)
	return err
}

func (c *AdodbConn) Begin() (driver.Tx, error) {
	_, err := oleutil.CallMethod(c.db, "BeginTrans")
	if err != nil {
		return nil, err
	}
	return &AdodbTx{c}, nil
}

func (d *AdodbDriver) Open(dsn string) (driver.Conn, error) {
	unknown, err := oleutil.CreateObject("ADODB.Connection")
	if err != nil {
		return nil, err
	}
	db, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.CallMethod(db, "Open", dsn)
	if err != nil {
		return nil, err
	}
	return &AdodbConn{db}, nil
}

func (c *AdodbConn) Close() error {
	_, err := oleutil.CallMethod(c.db, "Close")
	if err != nil {
		return err
	}
	c.db = nil
	return nil
}

type AdodbStmt struct {
	c  *AdodbConn
	s  *ole.IDispatch
	ps *ole.IDispatch
	b  []string
}

func (c *AdodbConn) Prepare(query string) (driver.Stmt, error) {
	unknown, err := oleutil.CreateObject("ADODB.Command")
	if err != nil {
		return nil, err
	}
	s, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "ActiveConnection", c.db)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandText", query)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandType", 1)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "Prepared", true)
	if err != nil {
		return nil, err
	}
	val, err := oleutil.GetProperty(s, "Parameters")
	if err != nil {
		return nil, err
	}
	return &AdodbStmt{c, s, val.ToIDispatch(), nil}, nil
}

func (s *AdodbStmt) Bind(bind []string) error {
	s.b = bind
	return nil
}

func (s *AdodbStmt) Close() error {
	s.s.Release()
	return nil
}

func (s *AdodbStmt) NumInput() int {
	if s.b != nil {
		return len(s.b)
	}
	_, err := oleutil.CallMethod(s.ps, "Refresh")
	if err != nil {
		return -1
	}
	val, err := oleutil.GetProperty(s.ps, "Count")
	if err != nil {
		return -1
	}
	c := int(val.Val)
	return c
}

func (s *AdodbStmt) bind(args []driver.Value) error {
	if s.b != nil {
		for i, v := range args {
			var b string = "?"
			if len(s.b) < i {
				b = s.b[i]
			}
			unknown, err := oleutil.CallMethod(s.s, "CreateParameter", b, 12, 1)
			if err != nil {
				return err
			}
			param := unknown.ToIDispatch()
			defer param.Release()
			_, err = oleutil.PutProperty(param, "Value", v)
			if err != nil {
				return err
			}
			_, err = oleutil.CallMethod(s.ps, "Append", param)
			if err != nil {
				return err
			}
		}
	} else {
		for i, v := range args {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(s.ps, "Item", &varval)
			if err != nil {
				return err
			}
			item := val.ToIDispatch()
			defer item.Release()
			_, err = oleutil.PutProperty(item, "Value", v)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *AdodbStmt) Query(args []driver.Value) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rc, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return &AdodbRows{s, rc.ToIDispatch(), -1, nil}, nil
}

func (s *AdodbStmt) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	_, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return driver.ResultNoRows, nil
}

type AdodbRows struct {
	s    *AdodbStmt
	rc   *ole.IDispatch
	nc   int
	cols []string
}

func (rc *AdodbRows) Close() error {
	_, err := oleutil.CallMethod(rc.rc, "Close")
	if err != nil {
		return err
	}
	return nil
}

func (rc *AdodbRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		unknown, err := oleutil.GetProperty(rc.rc, "Fields")
		if err != nil {
			return []string{}
		}
		fields := unknown.ToIDispatch()
		defer fields.Release()
		val, err := oleutil.GetProperty(fields, "Count")
		if err != nil {
			return []string{}
		}
		rc.nc = int(val.Val)
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(fields, "Item", &varval)
			if err != nil {
				return []string{}
			}
			item := val.ToIDispatch()
			if err != nil {
				return []string{}
			}
			name, err := oleutil.GetProperty(item, "Name")
			if err != nil {
				return []string{}
			}
			rc.cols[i] = name.ToString()
			item.Release()
		}
	}
	return rc.cols
}

func (rc *AdodbRows) Next(dest []driver.Value) error {
	unknown, err := oleutil.GetProperty(rc.rc, "EOF")
	if err != nil {
		return io.EOF
	}
	if unknown.Val != 0 {
		return io.EOF
	}
	unknown, err = oleutil.GetProperty(rc.rc, "Fields")
	if err != nil {
		return err
	}
	fields := unknown.ToIDispatch()
	defer fields.Release()
	for i := range dest {
		var varval ole.VARIANT
		varval.VT = ole.VT_I4
		varval.Val = int64(i)
		val, err := oleutil.CallMethod(fields, "Item", &varval)
		if err != nil {
			return err
		}
		field := val.ToIDispatch()
		defer field.Release()
		typ, err := oleutil.GetProperty(field, "Type")
		if err != nil {
			return err
		}
		val, err = oleutil.GetProperty(field, "Value")
		if err != nil {
			return err
		}
		field.Release()
		switch typ.Val {
		case 0: // ADEMPTY
			// TODO
		case 2: // ADSMALLINT
			dest[i] = int16(val.Val)
		case 3: // ADINTEGER
			dest[i] = int32(val.Val)
		case 4: // ADSINGLE
			dest[i] = float32(val.Val)
		case 5: // ADDOUBLE
			dest[i] = val.Val
		case 6: // ADCURRENCY
			dest[i] = float64(val.Val)
		case 7: // ADDATE
			// TODO
		case 8: // ADBSTR
			dest[i] = val.ToString()
		case 9: // ADIDISPATCH
			dest[i] = val.ToIDispatch()
		case 10: // ADERROR
			// TODO
		case 11: // ADBOOLEAN
			if val.Val != 0 {
				dest[i] = true
			} else {
				dest[i] = false
			}
		case 12: // ADVARIANT
			dest[i] = val
		case 13: // ADIUNKNOWN
			dest[i] = val.ToIUnknown()
		case 14: // ADDECIMAL
			dest[i] = float64(val.Val)
		case 16: // ADTINYINT
			dest[i] = int8(val.Val)
		case 17: // ADUNSIGNEDTINYINT
			dest[i] = uint8(val.Val)
		case 18: // ADUNSIGNEDSMALLINT
			dest[i] = uint16(val.Val)
		case 19: // ADUNSIGNEDINT
			dest[i] = uint32(val.Val)
		case 20: // ADBIGINT
			dest[i] = big.NewInt(val.Val)
		case 21: // ADUNSIGNEDBIGINT
			// TODO
		case 72: // ADGUID
			// TODO
		case 128: // ADBINARY
			sa := *(**ole.SAFEARRAY)(unsafe.Pointer(&val.Val))
			dest[i] = (*[1 << 30]byte)(unsafe.Pointer(uintptr(sa.PvData)))[0:sa.CbElements]
		case 129: // ADCHAR
			dest[i] = uint8(val.Val)
		case 130: // ADWCHAR
			dest[i] = uint16(val.Val)
		case 131: // ADNUMERIC
			dest[i] = val.Val
		case 132: // ADUSERDEFINED
			dest[i] = uintptr(val.Val)
		case 133: // ADDBDATE
			dest[i] = time.Unix(0, val.Val).UTC()
		case 134: // ADDBTIME
			dest[i] = time.Unix(0, val.Val).UTC()
		case 135: // ADDBTIMESTAMP
			dest[i] = time.Unix(0, val.Val).UTC()
		case 136: // ADCHAPTER
			dest[i] = val.ToString()
		case 200: // ADVARCHAR
			dest[i] = val.ToString()
		case 201: // ADLONGVARCHAR
			dest[i] = val.ToString()
		case 202: // ADVARWCHAR
			dest[i] = val.ToString()
		case 203: // ADLONGVARWCHAR
			dest[i] = val.ToString()
		case 204: // ADVARBINARY
			// TODO
		case 205: // ADLONGVARBINARY
			// TODO
		}
	}
	_, err = oleutil.CallMethod(rc.rc, "MoveNext")
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>sorry we have to remove import errors as well<commit_after>package adodb

import (
	"database/sql"
	"database/sql/driver"
	"github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
	"io"
	"math/big"
	"time"
	"unsafe"
)

func init() {
	ole.CoInitialize(0)
	sql.Register("adodb", &AdodbDriver{})
}

type AdodbDriver struct {

}

type AdodbConn struct {
	db *ole.IDispatch
}

type AdodbTx struct {
	c *AdodbConn
}

func (tx *AdodbTx) Commit() error {
	_, err := oleutil.CallMethod(tx.c.db, "CommitTrans")
	if err != nil {
		return err
	}
	return nil
}

func (tx *AdodbTx) Rollback() error {
	_, err := oleutil.CallMethod(tx.c.db, "Rollback")
	if err != nil {
		return err
	}
	return nil
}

func (c *AdodbConn) exec(cmd string) error {
	_, err := oleutil.CallMethod(c.db, "Execute", cmd)
	return err
}

func (c *AdodbConn) Begin() (driver.Tx, error) {
	_, err := oleutil.CallMethod(c.db, "BeginTrans")
	if err != nil {
		return nil, err
	}
	return &AdodbTx{c}, nil
}

func (d *AdodbDriver) Open(dsn string) (driver.Conn, error) {
	unknown, err := oleutil.CreateObject("ADODB.Connection")
	if err != nil {
		return nil, err
	}
	db, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.CallMethod(db, "Open", dsn)
	if err != nil {
		return nil, err
	}
	return &AdodbConn{db}, nil
}

func (c *AdodbConn) Close() error {
	_, err := oleutil.CallMethod(c.db, "Close")
	if err != nil {
		return err
	}
	c.db = nil
	return nil
}

type AdodbStmt struct {
	c  *AdodbConn
	s  *ole.IDispatch
	ps *ole.IDispatch
	b  []string
}

func (c *AdodbConn) Prepare(query string) (driver.Stmt, error) {
	unknown, err := oleutil.CreateObject("ADODB.Command")
	if err != nil {
		return nil, err
	}
	s, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "ActiveConnection", c.db)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandText", query)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandType", 1)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "Prepared", true)
	if err != nil {
		return nil, err
	}
	val, err := oleutil.GetProperty(s, "Parameters")
	if err != nil {
		return nil, err
	}
	return &AdodbStmt{c, s, val.ToIDispatch(), nil}, nil
}

func (s *AdodbStmt) Bind(bind []string) error {
	s.b = bind
	return nil
}

func (s *AdodbStmt) Close() error {
	s.s.Release()
	return nil
}

func (s *AdodbStmt) NumInput() int {
	if s.b != nil {
		return len(s.b)
	}
	_, err := oleutil.CallMethod(s.ps, "Refresh")
	if err != nil {
		return -1
	}
	val, err := oleutil.GetProperty(s.ps, "Count")
	if err != nil {
		return -1
	}
	c := int(val.Val)
	return c
}

func (s *AdodbStmt) bind(args []driver.Value) error {
	if s.b != nil {
		for i, v := range args {
			var b string = "?"
			if len(s.b) < i {
				b = s.b[i]
			}
			unknown, err := oleutil.CallMethod(s.s, "CreateParameter", b, 12, 1)
			if err != nil {
				return err
			}
			param := unknown.ToIDispatch()
			defer param.Release()
			_, err = oleutil.PutProperty(param, "Value", v)
			if err != nil {
				return err
			}
			_, err = oleutil.CallMethod(s.ps, "Append", param)
			if err != nil {
				return err
			}
		}
	} else {
		for i, v := range args {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(s.ps, "Item", &varval)
			if err != nil {
				return err
			}
			item := val.ToIDispatch()
			defer item.Release()
			_, err = oleutil.PutProperty(item, "Value", v)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *AdodbStmt) Query(args []driver.Value) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rc, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return &AdodbRows{s, rc.ToIDispatch(), -1, nil}, nil
}

func (s *AdodbStmt) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	_, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return driver.ResultNoRows, nil
}

type AdodbRows struct {
	s    *AdodbStmt
	rc   *ole.IDispatch
	nc   int
	cols []string
}

func (rc *AdodbRows) Close() error {
	_, err := oleutil.CallMethod(rc.rc, "Close")
	if err != nil {
		return err
	}
	return nil
}

func (rc *AdodbRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		unknown, err := oleutil.GetProperty(rc.rc, "Fields")
		if err != nil {
			return []string{}
		}
		fields := unknown.ToIDispatch()
		defer fields.Release()
		val, err := oleutil.GetProperty(fields, "Count")
		if err != nil {
			return []string{}
		}
		rc.nc = int(val.Val)
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(fields, "Item", &varval)
			if err != nil {
				return []string{}
			}
			item := val.ToIDispatch()
			if err != nil {
				return []string{}
			}
			name, err := oleutil.GetProperty(item, "Name")
			if err != nil {
				return []string{}
			}
			rc.cols[i] = name.ToString()
			item.Release()
		}
	}
	return rc.cols
}

func (rc *AdodbRows) Next(dest []driver.Value) error {
	unknown, err := oleutil.GetProperty(rc.rc, "EOF")
	if err != nil {
		return io.EOF
	}
	if unknown.Val != 0 {
		return io.EOF
	}
	unknown, err = oleutil.GetProperty(rc.rc, "Fields")
	if err != nil {
		return err
	}
	fields := unknown.ToIDispatch()
	defer fields.Release()
	for i := range dest {
		var varval ole.VARIANT
		varval.VT = ole.VT_I4
		varval.Val = int64(i)
		val, err := oleutil.CallMethod(fields, "Item", &varval)
		if err != nil {
			return err
		}
		field := val.ToIDispatch()
		defer field.Release()
		typ, err := oleutil.GetProperty(field, "Type")
		if err != nil {
			return err
		}
		val, err = oleutil.GetProperty(field, "Value")
		if err != nil {
			return err
		}
		field.Release()
		switch typ.Val {
		case 0: // ADEMPTY
			// TODO
		case 2: // ADSMALLINT
			dest[i] = int16(val.Val)
		case 3: // ADINTEGER
			dest[i] = int32(val.Val)
		case 4: // ADSINGLE
			dest[i] = float32(val.Val)
		case 5: // ADDOUBLE
			dest[i] = val.Val
		case 6: // ADCURRENCY
			dest[i] = float64(val.Val)
		case 7: // ADDATE
			// TODO
		case 8: // ADBSTR
			dest[i] = val.ToString()
		case 9: // ADIDISPATCH
			dest[i] = val.ToIDispatch()
		case 10: // ADERROR
			// TODO
		case 11: // ADBOOLEAN
			if val.Val != 0 {
				dest[i] = true
			} else {
				dest[i] = false
			}
		case 12: // ADVARIANT
			dest[i] = val
		case 13: // ADIUNKNOWN
			dest[i] = val.ToIUnknown()
		case 14: // ADDECIMAL
			dest[i] = float64(val.Val)
		case 16: // ADTINYINT
			dest[i] = int8(val.Val)
		case 17: // ADUNSIGNEDTINYINT
			dest[i] = uint8(val.Val)
		case 18: // ADUNSIGNEDSMALLINT
			dest[i] = uint16(val.Val)
		case 19: // ADUNSIGNEDINT
			dest[i] = uint32(val.Val)
		case 20: // ADBIGINT
			dest[i] = big.NewInt(val.Val)
		case 21: // ADUNSIGNEDBIGINT
			// TODO
		case 72: // ADGUID
			// TODO
		case 128: // ADBINARY
			sa := *(**ole.SAFEARRAY)(unsafe.Pointer(&val.Val))
			dest[i] = (*[1 << 30]byte)(unsafe.Pointer(uintptr(sa.PvData)))[0:sa.CbElements]
		case 129: // ADCHAR
			dest[i] = uint8(val.Val)
		case 130: // ADWCHAR
			dest[i] = uint16(val.Val)
		case 131: // ADNUMERIC
			dest[i] = val.Val
		case 132: // ADUSERDEFINED
			dest[i] = uintptr(val.Val)
		case 133: // ADDBDATE
			dest[i] = time.Unix(0, val.Val).UTC()
		case 134: // ADDBTIME
			dest[i] = time.Unix(0, val.Val).UTC()
		case 135: // ADDBTIMESTAMP
			dest[i] = time.Unix(0, val.Val).UTC()
		case 136: // ADCHAPTER
			dest[i] = val.ToString()
		case 200: // ADVARCHAR
			dest[i] = val.ToString()
		case 201: // ADLONGVARCHAR
			dest[i] = val.ToString()
		case 202: // ADVARWCHAR
			dest[i] = val.ToString()
		case 203: // ADLONGVARWCHAR
			dest[i] = val.ToString()
		case 204: // ADVARBINARY
			// TODO
		case 205: // ADLONGVARBINARY
			// TODO
		}
	}
	_, err = oleutil.CallMethod(rc.rc, "MoveNext")
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>