<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"testing"

	"github.com/bazelbuild/bazelisk/core"
	"github.com/bazelbuild/bazelisk/httputil"
	"github.com/bazelbuild/bazelisk/repositories"
	"github.com/bazelbuild/bazelisk/versions"
)

var (
	transport = &fakeTransport{responses: make(map[string]*http.Response)}
	tmpDir    = ""
)

func init() {
	httputil.DefaultTransport = transport
}

func TestMain(m *testing.M) {
	var err error
	tmpDir, err = ioutil.TempDir("", "version_test")
	if err != nil {
		log.Fatal(err)
	}
	defer os.RemoveAll(tmpDir)

	code := m.Run()
	os.Exit(code)
}

func TestResolveLatestRcVersion(t *testing.T) {
	s := setUp(t)
	s.AddVersion("4.0.0", false)
	s.AddVersion("10.0.0", false)
	s.AddVersion("11.0.0", true)
	s.AddVersion("11.11.0", false, 1, 2)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(nil, gcs, nil, nil, nil, false)
	version, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "last_rc")

	if err != nil {
		t.Fatalf("Version resolution failed unexpectedly: %v", err)
	}
	expectedRC := "11.11.0rc2"
	if version != expectedRC {
		t.Fatalf("Expected version %s, but got %s", expectedRC, version)
	}
}

func TestResolveLatestRcVersion_WithFullRelease(t *testing.T) {
	s := setUp(t)
	s.AddVersion("4.0.0", true, 1, 2, 3)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(nil, gcs, nil, nil, nil, false)
	version, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "last_rc")

	if err != nil {
		t.Fatalf("Version resolution failed unexpectedly: %v", err)
	}
	expectedRC := "4.0.0rc3"
	if version != expectedRC {
		t.Fatalf("Expected version %s, but got %s", expectedRC, version)
	}
}

func TestResolveLatestVersion_TwoLatestVersionsDoNotHaveAReleaseYet(t *testing.T) {
	s := setUp(t)
	s.AddVersion("4.0.0", true)
	s.AddVersion("5.0.0", false)
	s.AddVersion("6.0.0", false)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(gcs, nil, nil, nil, nil, false)
	version, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "latest")

	if err != nil {
		t.Fatalf("Version resolution failed unexpectedly: %v", err)
	}
	expectedVersion := "4.0.0"
	if version != expectedVersion {
		t.Fatalf("Expected version %s, but got %s", expectedVersion, version)
	}
}

func TestResolveLatestVersion_FilterReleaseCandidates(t *testing.T) {
	s := setUp(t)
	s.AddVersion("3.0.0", true)
	s.AddVersion("4.0.0", false)
	s.AddVersion("5.0.0", false)
	s.AddVersion("6.0.0", true)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(gcs, nil, nil, nil, nil, false)
	version, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "latest-1")

	if err != nil {
		t.Fatalf("Version resolution failed unexpectedly: %v", err)
	}
	expectedVersion := "3.0.0"
	if version != expectedVersion {
		t.Fatalf("Expected version %s, but got %s", expectedVersion, version)
	}
}

func TestResolveLatestVersion_ShouldFailIfNotEnoughReleases(t *testing.T) {
	s := setUp(t)
	s.AddVersion("3.0.0", true)
	s.AddVersion("4.0.0", false)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(gcs, nil, nil, nil, nil, false)
	_, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "latest-1")

	if err == nil {
		t.Fatal("Expected ResolveVersion() to fail.")
	}
	expectedError := "unable to determine latest version: requested 2 latest releases, but only found 1"
	if err.Error() != expectedError {
		t.Fatalf("Expected error message %q, but got '%v'", expectedError, err)
	}
}

func TestResolveLatestVersion_GCSIsDown(t *testing.T) {
	setUp(t).WithError().Finish()

	transport.AddResponse("https://www.googleapis.com/storage/v1/b/bazel/o?delimiter=/", 500, "")

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(gcs, nil, nil, nil, nil, false)
	_, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "latest")

	if err == nil {
		t.Fatal("Expected resolveLatestVersion() to fail.")
	}
	expectedPrefix := "unable to determine latest version: could not list Bazel versions in GCS bucket"
	if !strings.HasPrefix(err.Error(), expectedPrefix) {
		t.Fatalf("Expected error message that starts with %q, but got '%v'", expectedPrefix, err)
	}
}

func TestResolveLatestVersion_GitHubIsDown(t *testing.T) {
	transport.AddResponse("https://api.github.com/repos/bazelbuild/bazel/releases", 500, "")

	gh := repositories.CreateGitHubRepo("test_token")
	repos := core.CreateRepositories(nil, nil, gh, nil, nil, false)

	_, _, err := repos.ResolveVersion(tmpDir, "some_fork", "latest")

	if err == nil {
		t.Fatal("Expected resolveLatestVersion() to fail.")
	}
	expectedPrefix := "unable to determine latest version: unable to dermine 'some_fork' releases: could not download list of Bazel releases from github.com/some_fork"
	if !strings.HasPrefix(err.Error(), expectedPrefix) {
		t.Fatalf("Expected error message that starts with %q, but got '%v'", expectedPrefix, err)
	}
}

func TestAcceptRollingReleaseName(t *testing.T) {
	gh := repositories.CreateGitHubRepo("test_token")
	repos := core.CreateRepositories(nil, nil, nil, nil, gh, false)

	version := "10.0.0-pre.20201103.4"

	resolvedVersion, _, err := repos.ResolveVersion(tmpDir, "", version)

	if err != nil {
		t.Fatalf("ResolveVersion(%q, \"\", %q): expected no error, but got %v", tmpDir, version, err)
	}

	if resolvedVersion != version {
		t.Fatalf("ResolveVersion(%q, \"\", %q) = %v, but expected %v", tmpDir, version, resolvedVersion, version)
	}
}

func TestResolveLatestRollingRelease(t *testing.T) {
	text := `
	[
	  {
		"tag_name": "4.0.0",
		"prerelease": false
	  },
	  {
		"tag_name": "5.0.0-pre.20210319",
		"prerelease": true
	  },
	  {
		"tag_name": "5.0.0-pre.20210322.4",
		"prerelease": true
	  },
	  {
		"tag_name": "5.0.0",
		"prerelease": false
	  }
	]
	`

	transport.AddResponse("https://api.github.com/repos/bazelbuild/bazel/releases", 200, text)

	gh := repositories.CreateGitHubRepo("test_token")
	repos := core.CreateRepositories(nil, nil, nil, nil, gh, false)

	version, _, err := repos.ResolveVersion(tmpDir, "", "rolling")

	if err != nil {
		t.Fatalf("ResolveVersion(%q, \"\", \"rolling\"): expected no error, but got %v", tmpDir, err)
	}

	want := "5.0.0-pre.20210322.4"
	if version != want {
		t.Fatalf("ResolveVersion(%q, \"\", \"rolling\") = %v, but expected %v", tmpDir, version, want)
	}
}

type gcsSetup struct {
	baseURL         string
	versionPrefixes []string
	status          int
	test            *testing.T
}

func (g *gcsSetup) AddVersion(version string, hasRelease bool, rcs ...int) *gcsSetup {
	g.versionPrefixes = append(g.versionPrefixes, fmt.Sprintf("%s/", version))
	prefixes := make([]string, 0)
	for _, rc := range rcs {
		prefix := fmt.Sprintf("%s/rc%d/", version, rc)
		prefixes = append(prefixes, prefix)
		g.addURL(prefix, false)
	}

	// The /release/ URLs have to exist, even if there is no release. In this case GCS returns no items, though.
	releasePrefix := fmt.Sprintf("%s/release/", version)
	g.addURL(releasePrefix, hasRelease)
	if hasRelease {
		prefixes = append(prefixes, releasePrefix)
	}

	g.addURL(fmt.Sprintf("%s/", version), false, prefixes...)
	return g
}

func (g *gcsSetup) addURL(prefix string, containsItem bool, childPrefixes ...string) {
	items := make([]interface{}, 0)
	if containsItem {
		items = append(items, "this_is_a_release")
	}
	resp := buildGCSResponseOrFail(g.test, childPrefixes, items)
	transport.AddResponse(fmt.Sprintf("%s&prefix=%s", g.baseURL, prefix), 200, resp)
}

func setUp(t *testing.T) *gcsSetup {
	return &gcsSetup{
		baseURL:         "https://www.googleapis.com/storage/v1/b/bazel/o?delimiter=/",
		status:          200,
		versionPrefixes: make([]string, 0),
		test:            t,
	}
}

func (g *gcsSetup) WithError() *gcsSetup {
	g.status = 500
	return g
}

func (g *gcsSetup) Finish() {
	// TODO: sort and deduplicate versionPrefixes
	listBody := buildGCSResponseOrFail(g.test, g.versionPrefixes, []interface{}{})
	transport.AddResponse(g.baseURL, g.status, listBody)
}

type fakeTransport struct {
	responses map[string]*http.Response
}

func (ft *fakeTransport) AddResponse(url string, status int, body string) {
	ft.responses[url] = ft.createResponse(status, body)
}

func (ft *fakeTransport) createResponse(status int, body string) *http.Response {
	return &http.Response{
		StatusCode: status,
		Body:       ioutil.NopCloser(bytes.NewBufferString(body)),
	}
}

func (ft *fakeTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if resp, ok := ft.responses[req.URL.String()]; ok {
		return resp, nil
	}
	return ft.createResponse(http.StatusNotFound, ""), nil
}

func buildGCSResponseOrFail(t *testing.T, prefixes []string, items []interface{}) string {
	r := &repositories.GcsListResponse{
		Prefixes: prefixes,
		Items:    items,
	}
	byteValue, err := json.Marshal(r)
	if err != nil {
		t.Fatalf("Could not build GCS json response: %v", err)
		return ""
	}
	return string(byteValue)
}
<commit_msg>Apply style review feedback (part 2) (#233)<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"testing"

	"github.com/bazelbuild/bazelisk/core"
	"github.com/bazelbuild/bazelisk/httputil"
	"github.com/bazelbuild/bazelisk/repositories"
	"github.com/bazelbuild/bazelisk/versions"
)

const (
	rollingReleaseIdentifier = "rolling"
)

var (
	transport = &fakeTransport{responses: make(map[string]*http.Response)}
	tmpDir    = ""
)

func init() {
	httputil.DefaultTransport = transport
}

func TestMain(m *testing.M) {
	var err error
	tmpDir, err = ioutil.TempDir("", "version_test")
	if err != nil {
		log.Fatal(err)
	}
	defer os.RemoveAll(tmpDir)

	code := m.Run()
	os.Exit(code)
}

func TestResolveLatestRcVersion(t *testing.T) {
	s := setUp(t)
	s.AddVersion("4.0.0", false)
	s.AddVersion("10.0.0", false)
	s.AddVersion("11.0.0", true)
	s.AddVersion("11.11.0", false, 1, 2)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(nil, gcs, nil, nil, nil, false)
	version, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "last_rc")

	if err != nil {
		t.Fatalf("Version resolution failed unexpectedly: %v", err)
	}
	expectedRC := "11.11.0rc2"
	if version != expectedRC {
		t.Fatalf("Expected version %s, but got %s", expectedRC, version)
	}
}

func TestResolveLatestRcVersion_WithFullRelease(t *testing.T) {
	s := setUp(t)
	s.AddVersion("4.0.0", true, 1, 2, 3)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(nil, gcs, nil, nil, nil, false)
	version, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "last_rc")

	if err != nil {
		t.Fatalf("Version resolution failed unexpectedly: %v", err)
	}
	expectedRC := "4.0.0rc3"
	if version != expectedRC {
		t.Fatalf("Expected version %s, but got %s", expectedRC, version)
	}
}

func TestResolveLatestVersion_TwoLatestVersionsDoNotHaveAReleaseYet(t *testing.T) {
	s := setUp(t)
	s.AddVersion("4.0.0", true)
	s.AddVersion("5.0.0", false)
	s.AddVersion("6.0.0", false)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(gcs, nil, nil, nil, nil, false)
	version, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "latest")

	if err != nil {
		t.Fatalf("Version resolution failed unexpectedly: %v", err)
	}
	expectedVersion := "4.0.0"
	if version != expectedVersion {
		t.Fatalf("Expected version %s, but got %s", expectedVersion, version)
	}
}

func TestResolveLatestVersion_FilterReleaseCandidates(t *testing.T) {
	s := setUp(t)
	s.AddVersion("3.0.0", true)
	s.AddVersion("4.0.0", false)
	s.AddVersion("5.0.0", false)
	s.AddVersion("6.0.0", true)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(gcs, nil, nil, nil, nil, false)
	version, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "latest-1")

	if err != nil {
		t.Fatalf("Version resolution failed unexpectedly: %v", err)
	}
	expectedVersion := "3.0.0"
	if version != expectedVersion {
		t.Fatalf("Expected version %s, but got %s", expectedVersion, version)
	}
}

func TestResolveLatestVersion_ShouldFailIfNotEnoughReleases(t *testing.T) {
	s := setUp(t)
	s.AddVersion("3.0.0", true)
	s.AddVersion("4.0.0", false)
	s.Finish()

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(gcs, nil, nil, nil, nil, false)
	_, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "latest-1")

	if err == nil {
		t.Fatal("Expected ResolveVersion() to fail.")
	}
	expectedError := "unable to determine latest version: requested 2 latest releases, but only found 1"
	if err.Error() != expectedError {
		t.Fatalf("Expected error message %q, but got '%v'", expectedError, err)
	}
}

func TestResolveLatestVersion_GCSIsDown(t *testing.T) {
	setUp(t).WithError().Finish()

	transport.AddResponse("https://www.googleapis.com/storage/v1/b/bazel/o?delimiter=/", 500, "")

	gcs := &repositories.GCSRepo{}
	repos := core.CreateRepositories(gcs, nil, nil, nil, nil, false)
	_, _, err := repos.ResolveVersion(tmpDir, versions.BazelUpstream, "latest")

	if err == nil {
		t.Fatal("Expected resolveLatestVersion() to fail.")
	}
	expectedPrefix := "unable to determine latest version: could not list Bazel versions in GCS bucket"
	if !strings.HasPrefix(err.Error(), expectedPrefix) {
		t.Fatalf("Expected error message that starts with %q, but got '%v'", expectedPrefix, err)
	}
}

func TestResolveLatestVersion_GitHubIsDown(t *testing.T) {
	transport.AddResponse("https://api.github.com/repos/bazelbuild/bazel/releases", 500, "")

	gh := repositories.CreateGitHubRepo("test_token")
	repos := core.CreateRepositories(nil, nil, gh, nil, nil, false)

	_, _, err := repos.ResolveVersion(tmpDir, "some_fork", "latest")

	if err == nil {
		t.Fatal("Expected resolveLatestVersion() to fail.")
	}
	expectedPrefix := "unable to determine latest version: unable to dermine 'some_fork' releases: could not download list of Bazel releases from github.com/some_fork"
	if !strings.HasPrefix(err.Error(), expectedPrefix) {
		t.Fatalf("Expected error message that starts with %q, but got '%v'", expectedPrefix, err)
	}
}

func TestAcceptRollingReleaseName(t *testing.T) {
	gh := repositories.CreateGitHubRepo("test_token")
	repos := core.CreateRepositories(nil, nil, nil, nil, gh, false)

	version := "10.0.0-pre.20201103.4"

	resolvedVersion, _, err := repos.ResolveVersion(tmpDir, "", version)

	if err != nil {
		t.Fatalf("ResolveVersion(%q, \"\", %q): expected no error, but got %v", tmpDir, version, err)
	}

	if resolvedVersion != version {
		t.Fatalf("ResolveVersion(%q, \"\", %q) = %v, but expected %v", tmpDir, version, resolvedVersion, version)
	}
}

func TestResolveLatestRollingRelease(t *testing.T) {
	text := `
	[
	  {
		"tag_name": "4.0.0",
		"prerelease": false
	  },
	  {
		"tag_name": "5.0.0-pre.20210319",
		"prerelease": true
	  },
	  {
		"tag_name": "5.0.0-pre.20210322.4",
		"prerelease": true
	  },
	  {
		"tag_name": "5.0.0",
		"prerelease": false
	  }
	]
	`

	transport.AddResponse("https://api.github.com/repos/bazelbuild/bazel/releases", 200, text)

	gh := repositories.CreateGitHubRepo("test_token")
	repos := core.CreateRepositories(nil, nil, nil, nil, gh, false)

	version, _, err := repos.ResolveVersion(tmpDir, "", rollingReleaseIdentifier)

	if err != nil {
		t.Fatalf("ResolveVersion(%q, \"\", %q): expected no error, but got %v", tmpDir, rollingReleaseIdentifier, err)
	}

	want := "5.0.0-pre.20210322.4"
	if version != want {
		t.Fatalf("ResolveVersion(%q, \"\", %q) = %v, but expected %v", tmpDir, rollingReleaseIdentifier, version, want)
	}
}

type gcsSetup struct {
	baseURL         string
	versionPrefixes []string
	status          int
	test            *testing.T
}

func (g *gcsSetup) AddVersion(version string, hasRelease bool, rcs ...int) *gcsSetup {
	g.versionPrefixes = append(g.versionPrefixes, fmt.Sprintf("%s/", version))
	prefixes := make([]string, 0)
	for _, rc := range rcs {
		prefix := fmt.Sprintf("%s/rc%d/", version, rc)
		prefixes = append(prefixes, prefix)
		g.addURL(prefix, false)
	}

	// The /release/ URLs have to exist, even if there is no release. In this case GCS returns no items, though.
	releasePrefix := fmt.Sprintf("%s/release/", version)
	g.addURL(releasePrefix, hasRelease)
	if hasRelease {
		prefixes = append(prefixes, releasePrefix)
	}

	g.addURL(fmt.Sprintf("%s/", version), false, prefixes...)
	return g
}

func (g *gcsSetup) addURL(prefix string, containsItem bool, childPrefixes ...string) {
	items := make([]interface{}, 0)
	if containsItem {
		items = append(items, "this_is_a_release")
	}
	resp := buildGCSResponseOrFail(g.test, childPrefixes, items)
	transport.AddResponse(fmt.Sprintf("%s&prefix=%s", g.baseURL, prefix), 200, resp)
}

func setUp(t *testing.T) *gcsSetup {
	return &gcsSetup{
		baseURL:         "https://www.googleapis.com/storage/v1/b/bazel/o?delimiter=/",
		status:          200,
		versionPrefixes: make([]string, 0),
		test:            t,
	}
}

func (g *gcsSetup) WithError() *gcsSetup {
	g.status = 500
	return g
}

func (g *gcsSetup) Finish() {
	// TODO: sort and deduplicate versionPrefixes
	listBody := buildGCSResponseOrFail(g.test, g.versionPrefixes, []interface{}{})
	transport.AddResponse(g.baseURL, g.status, listBody)
}

type fakeTransport struct {
	responses map[string]*http.Response
}

func (ft *fakeTransport) AddResponse(url string, status int, body string) {
	ft.responses[url] = ft.createResponse(status, body)
}

func (ft *fakeTransport) createResponse(status int, body string) *http.Response {
	return &http.Response{
		StatusCode: status,
		Body:       ioutil.NopCloser(bytes.NewBufferString(body)),
	}
}

func (ft *fakeTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if resp, ok := ft.responses[req.URL.String()]; ok {
		return resp, nil
	}
	return ft.createResponse(http.StatusNotFound, ""), nil
}

func buildGCSResponseOrFail(t *testing.T, prefixes []string, items []interface{}) string {
	r := &repositories.GcsListResponse{
		Prefixes: prefixes,
		Items:    items,
	}
	byteValue, err := json.Marshal(r)
	if err != nil {
		t.Fatalf("Could not build GCS json response: %v", err)
		return ""
	}
	return string(byteValue)
}
<|endoftext|>


<commit_before>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`
const start = `start`
const stop = `stop`

var eventsDemand = regexp.MustCompile(`^events`)
var logsDemand = regexp.MustCompile(`^logs (.+)`)
var statsDemand = regexp.MustCompile(`^stats (.+) (.+)`)
var busWebsocketRequest = regexp.MustCompile(`bus`)
var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var statsWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/stats`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)

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

func readUntilClose(user *auth.User, ws *websocket.Conn, name string) bool {
	messageType, _, err := ws.ReadMessage()

	if messageType == websocket.CloseMessage {
		return true
	}

	if err != nil {
		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
			log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
		}

		return true
	}

	return false
}

func readContent(user *auth.User, ws *websocket.Conn, name string, done chan<- int, content chan<- []byte) {
	for {
		messageType, message, err := ws.ReadMessage()

		if messageType == websocket.CloseMessage {
			close(done)
			return
		}

		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
				log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
			}

			close(done)
			return
		}

		content <- message
	}
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth))
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))

		defer ws.Close()
		return nil, nil, err
	}

	return ws, user, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	logs, err := docker.ContainerLogs(ctx, string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}
	defer logs.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				return

			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize {
					if err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
						log.Printf(`[%s] Error while writing to logs socket: %v`, user.Username, err)
						return
					}
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			break

		default:
			if readUntilClose(user, ws, `logs`) {
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	filtersArgs := filters.NewArgs()
	if labelFilters(&filtersArgs, user, nil) != nil {
		log.Printf(`[%s] Error while defining label filters: %v`, user.Username, err)
		return
	}
	if eventFilters(&filtersArgs) != nil {
		log.Printf(`[%s] Error while defining event filters: %v`, user.Username, err)
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	messages, errors := docker.Events(ctx, types.EventsOptions{Filters: filtersArgs})

	go func() {
		defer cancel()

		for {
			select {
			case <-ctx.Done():
				return

			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Printf(`[%s] Error while marshalling event: %v`, user.Username, err)
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Printf(`[%s] Error while writing to events socket: %v`, user.Username, err)
					return
				}

			case err := <-errors:
				log.Printf(`[%s] Error while reading events: %v`, user.Username, err)
				return
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			break

		default:
			if readUntilClose(user, ws, `events`) {
				return
			}
		}
	}
}

func statsWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	stats, err := docker.ContainerStats(ctx, string(containerID), true)
	if err != nil {
		log.Print(err)
		return
	}
	defer stats.Body.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(stats.Body)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				log.Printf(`[%s] Stats context is over for writing`, user.Username)
				return

			default:
				if err = ws.WriteMessage(websocket.TextMessage, scanner.Bytes()); err != nil {
					log.Printf(`[%s] Error while writing to stats socket: %v`, user.Username, err)
					return
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			log.Printf(`[%s] Stats context is over for reading`, user.Username)
			return

		default:
			if readUntilClose(user, ws, `stats`) {
				return
			}
		}
	}
}

func streamStats(ctx *context.Context, user *auth.User, containerID string, output chan<- []byte) {
	stats, err := docker.ContainerStats(ctx, containerID, true)
	if err != nil {
		log.Printf(`[%s] Error while streaming stats: %v`, user.Username, err)
		return
	}
	defer stats.Body.Close()

	scanner := bufio.NewScanner(stats.Body)
	for scanner.Scan() {
		output <- scanner.Bytes()
	}
}

func busWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	done := make(chan int)

	output := make(chan []byte)
	defer close(output)

	input := make(chan []byte)
	defer close(input)

	go readContent(user, ws, `streaming`, done, input)
	log.Printf(`[%s] Streaming started`, user.Username)

	var statsCancelFunc context.CancelFunc

	for {
		select {
		case <-done:
			log.Printf(`[%s] Streaming ended`, user.Username)
			return

		case inputBytes := <-input:
			if eventsDemand.Match(inputBytes) {
				log.Printf(`[%s] Streaming events`, user.Username)
			} else if logsDemand.Match(inputBytes) {
				log.Printf(`[%s] Streaming logs for %s`, user.Username, logsDemand.FindSubmatch(inputBytes)[1])
			} else if statsDemand.Match(inputBytes) {
				containerID := statsDemand.FindSubmatch(inputBytes)[1]
				action := string(statsDemand.FindSubmatch(inputBytes)[2])

				if action == stop && cancelFunc != nil {
					log.Printf(`[%s] Stopping stats stream`, user.Username)
					statsCancelFunc()
				} else if action == start {
					log.Printf(`[%s] Starting stats stream for %s`, user.Username, containerID)

					if statsCancelFunc != nil {
						log.Printf(`[%s] Cancelling previous stats stream`, user.Username)
						statsCancelFunc()
					}
					statsContext, statsCancelFunc := context.WithCancel(context.Background())
					defer statsCancelFunc()

					streamStats(statsContext, user, containerID, output)
				}
			}

		case outputBytes := <-output:
			if err = ws.WriteMessage(websocket.TextMessage, outputBytes); err != nil {
				log.Printf(`[%s] Error while writing to streaming: %v`, user.Username, err)
				close(done)
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	} else if statsWebsocketRequest.Match(urlPath) {
		statsWebsocketHandler(w, r, statsWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if busWebsocketRequest.Match(urlPath) {
		busWebsocketHandler(w, r)
	}
}
<commit_msg>Update websocket.go<commit_after>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`
const start = `start`
const stop = `stop`

var eventsDemand = regexp.MustCompile(`^events`)
var logsDemand = regexp.MustCompile(`^logs (.+)`)
var statsDemand = regexp.MustCompile(`^stats (.+) (.+)`)
var busWebsocketRequest = regexp.MustCompile(`bus`)
var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var statsWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/stats`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)

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

func readUntilClose(user *auth.User, ws *websocket.Conn, name string) bool {
	messageType, _, err := ws.ReadMessage()

	if messageType == websocket.CloseMessage {
		return true
	}

	if err != nil {
		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
			log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
		}

		return true
	}

	return false
}

func readContent(user *auth.User, ws *websocket.Conn, name string, done chan<- int, content chan<- []byte) {
	for {
		messageType, message, err := ws.ReadMessage()

		if messageType == websocket.CloseMessage {
			close(done)
			return
		}

		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
				log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
			}

			close(done)
			return
		}

		content <- message
	}
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth))
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))

		defer ws.Close()
		return nil, nil, err
	}

	return ws, user, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	logs, err := docker.ContainerLogs(ctx, string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}
	defer logs.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				return

			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize {
					if err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
						log.Printf(`[%s] Error while writing to logs socket: %v`, user.Username, err)
						return
					}
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			break

		default:
			if readUntilClose(user, ws, `logs`) {
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	filtersArgs := filters.NewArgs()
	if labelFilters(&filtersArgs, user, nil) != nil {
		log.Printf(`[%s] Error while defining label filters: %v`, user.Username, err)
		return
	}
	if eventFilters(&filtersArgs) != nil {
		log.Printf(`[%s] Error while defining event filters: %v`, user.Username, err)
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	messages, errors := docker.Events(ctx, types.EventsOptions{Filters: filtersArgs})

	go func() {
		defer cancel()

		for {
			select {
			case <-ctx.Done():
				return

			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Printf(`[%s] Error while marshalling event: %v`, user.Username, err)
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Printf(`[%s] Error while writing to events socket: %v`, user.Username, err)
					return
				}

			case err := <-errors:
				log.Printf(`[%s] Error while reading events: %v`, user.Username, err)
				return
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			break

		default:
			if readUntilClose(user, ws, `events`) {
				return
			}
		}
	}
}

func statsWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	stats, err := docker.ContainerStats(ctx, string(containerID), true)
	if err != nil {
		log.Print(err)
		return
	}
	defer stats.Body.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(stats.Body)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				log.Printf(`[%s] Stats context is over for writing`, user.Username)
				return

			default:
				if err = ws.WriteMessage(websocket.TextMessage, scanner.Bytes()); err != nil {
					log.Printf(`[%s] Error while writing to stats socket: %v`, user.Username, err)
					return
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			log.Printf(`[%s] Stats context is over for reading`, user.Username)
			return

		default:
			if readUntilClose(user, ws, `stats`) {
				return
			}
		}
	}
}

func streamStats(ctx context.Context, user *auth.User, containerID string, output chan<- []byte) {
	stats, err := docker.ContainerStats(ctx, containerID, true)
	if err != nil {
		log.Printf(`[%s] Error while streaming stats: %v`, user.Username, err)
		return
	}
	defer stats.Body.Close()

	scanner := bufio.NewScanner(stats.Body)
	for scanner.Scan() {
		output <- scanner.Bytes()
	}
}

func busWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	done := make(chan int)

	output := make(chan []byte)
	defer close(output)

	input := make(chan []byte)
	defer close(input)

	go readContent(user, ws, `streaming`, done, input)
	log.Printf(`[%s] Streaming started`, user.Username)

	var statsCancelFunc context.CancelFunc

	for {
		select {
		case <-done:
			log.Printf(`[%s] Streaming ended`, user.Username)
			return

		case inputBytes := <-input:
			if eventsDemand.Match(inputBytes) {
				log.Printf(`[%s] Streaming events`, user.Username)
			} else if logsDemand.Match(inputBytes) {
				log.Printf(`[%s] Streaming logs for %s`, user.Username, logsDemand.FindSubmatch(inputBytes)[1])
			} else if statsDemand.Match(inputBytes) {
				containerID := statsDemand.FindSubmatch(inputBytes)[1]
				action := string(statsDemand.FindSubmatch(inputBytes)[2])

				if action == stop && statsCancelFunc != nil {
					log.Printf(`[%s] Stopping stats stream`, user.Username)
					statsCancelFunc()
				} else if action == start {
					log.Printf(`[%s] Starting stats stream for %s`, user.Username, containerID)

					if statsCancelFunc != nil {
						log.Printf(`[%s] Cancelling previous stats stream`, user.Username)
						statsCancelFunc()
					}
					statsContext, statsCancelFunc := context.WithCancel(context.Background())
					defer statsCancelFunc()

					streamStats(statsContext, user, string(containerID), output)
				}
			}

		case outputBytes := <-output:
			if err = ws.WriteMessage(websocket.TextMessage, outputBytes); err != nil {
				log.Printf(`[%s] Error while writing to streaming: %v`, user.Username, err)
				close(done)
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	} else if statsWebsocketRequest.Match(urlPath) {
		statsWebsocketHandler(w, r, statsWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if busWebsocketRequest.Match(urlPath) {
		busWebsocketHandler(w, r)
	}
}
<|endoftext|>


<commit_before>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`

var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var statsWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/stats`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)

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

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth))
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))

		defer ws.Close()
		return nil, nil, err
	}

	return ws, user, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, _, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	logs, err := docker.ContainerLogs(context, string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}
	defer logs.Close()

	done := make(chan struct{})

	go func() {
		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				return

			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize {
					if err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
						log.Printf(`Error while writing to logs socket: %v`, err)
						cancel()
						return
					}
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return

		default:
			if messageType, content, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from logs socket: %v | %v | %v`, err, messageType, content)
				cancel()
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	filtersArgs := filters.NewArgs()
	if labelFilters(&filtersArgs, user, nil) != nil {
		log.Printf(`Error while defining label filters: %v`, err)
		return
	}
	if eventFilters(&filtersArgs) != nil {
		log.Printf(`Error while defining event filters: %v`, err)
		return
	}

	context := context.Background()
	defer context.Done()
	messages, errors := docker.Events(context, types.EventsOptions{Filters: filtersArgs})

	done := make(chan struct{})

	go func() {
		for {
			select {
			case <-done:
				return

			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Printf(`Error while marshalling event: %v`, err)
					close(done)
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Printf(`Error while writing to events socket: %v`, err)
					close(done)
					return
				}
				break

			case err := <-errors:
				log.Printf(`Error while reading events: %v`, err)
				close(done)
				return
			}
		}
	}()

	for {
		select {
		case <-done:
			return
		default:
			if _, _, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from events socket: %v`, err)
				close(done)
				return
			}
		}
	}
}

func statsWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, _, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	context := context.Background()
	defer context.Done()

	stats, err := docker.ContainerStats(context, string(containerID), true)
	if err != nil {
		log.Print(err)
		return
	}
	defer stats.Body.Close()

	done := make(chan struct{})

	go func() {
		scanner := bufio.NewScanner(stats.Body)
		for scanner.Scan() {
			select {
			case <-done:
				return

			default:
				if err = ws.WriteMessage(websocket.TextMessage, scanner.Bytes()); err != nil {
					log.Printf(`Error while writing to stats socket: %v`, err)
					close(done)
					return
				}
			}
		}
	}()

	for {
		select {
		case <-done:
			return

		default:
			if _, _, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from stats socket: %v`, err)
				close(done)
				return
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	} else if statsWebsocketRequest.Match(urlPath) {
		statsWebsocketHandler(w, r, statsWebsocketRequest.FindSubmatch(urlPath)[1])
	}
}
<commit_msg>Update websocket.go<commit_after>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`

var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var statsWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/stats`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)

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

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth))
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))

		defer ws.Close()
		return nil, nil, err
	}

	return ws, user, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, _, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	logs, err := docker.ContainerLogs(ctx, string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}
	defer logs.Close()

	go func() {
		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				return

			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize {
					if err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
						log.Printf(`Error while writing to logs socket: %v`, err)
						cancel()
						return
					}
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return

		default:
			if messageType, content, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from logs socket: %v | %v | %v`, err, messageType, content)
				cancel()
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	filtersArgs := filters.NewArgs()
	if labelFilters(&filtersArgs, user, nil) != nil {
		log.Printf(`Error while defining label filters: %v`, err)
		return
	}
	if eventFilters(&filtersArgs) != nil {
		log.Printf(`Error while defining event filters: %v`, err)
		return
	}

	context := context.Background()
	defer context.Done()
	messages, errors := docker.Events(context, types.EventsOptions{Filters: filtersArgs})

	done := make(chan struct{})

	go func() {
		for {
			select {
			case <-done:
				return

			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Printf(`Error while marshalling event: %v`, err)
					close(done)
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Printf(`Error while writing to events socket: %v`, err)
					close(done)
					return
				}
				break

			case err := <-errors:
				log.Printf(`Error while reading events: %v`, err)
				close(done)
				return
			}
		}
	}()

	for {
		select {
		case <-done:
			return
		default:
			if _, _, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from events socket: %v`, err)
				close(done)
				return
			}
		}
	}
}

func statsWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, _, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	context := context.Background()
	defer context.Done()

	stats, err := docker.ContainerStats(context, string(containerID), true)
	if err != nil {
		log.Print(err)
		return
	}
	defer stats.Body.Close()

	done := make(chan struct{})

	go func() {
		scanner := bufio.NewScanner(stats.Body)
		for scanner.Scan() {
			select {
			case <-done:
				return

			default:
				if err = ws.WriteMessage(websocket.TextMessage, scanner.Bytes()); err != nil {
					log.Printf(`Error while writing to stats socket: %v`, err)
					close(done)
					return
				}
			}
		}
	}()

	for {
		select {
		case <-done:
			return

		default:
			if _, _, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from stats socket: %v`, err)
				close(done)
				return
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	} else if statsWebsocketRequest.Match(urlPath) {
		statsWebsocketHandler(w, r, statsWebsocketRequest.FindSubmatch(urlPath)[1])
	}
}
<|endoftext|>


<commit_before>package spiders

// where i defined spider
import (
	"github.com/wcong/ants-go/ants/http"
	"log"
)

const (
	BASE_PARSE_NAME     = "base"
	SPIDERS_STATUS_INIT = iota
	SPIDERS_STATUS_RUNNING
	SPIDERS_STATUS_STOP
	SPIDERS_BASIC_COOKIE
)

/*
what a spider do
*	make start request
* 	define basic parse func
*/
type Spider struct {
	BeforeMethod  func()
	InitStartUrls func()
	AfterMethod   func()
	Name          string
	StartUrls     []string
	ExtData       map[string]interface{}
	ParseMap      map[string]func(response *http.Response) ([]*http.Request, error) // defined you own parse function to scrapy response,make sure it return []response
}

func (this *Spider) MakeStartRequests() []*http.Request {
	if this.InitStartUrls != nil {
		this.InitStartUrls()
	}
	startRequestSlice := make([]*http.Request, len(this.StartUrls))
	for index, url := range this.StartUrls {
		request, err := http.NewRequest("GET", url, this.Name, BASE_PARSE_NAME, nil, 0)
		if err != nil {
			log.Println(err)
			continue
		}
		startRequestSlice[index] = request
	}
	return startRequestSlice
}
<commit_msg>optimize skip list<commit_after>package main

import (
	"math/rand"
)

type IntKey struct {
	Key int
}

func NewIntKey(key int) *IntKey {
	return &IntKey{key}
}

func (this *IntKey) GetNum() int {
	return this.Key
}
func (this *IntKey) Compare(key Key) int8 {
	return (int8)(this.Key - key.GetNum())
}

/*
* skip list
**/

type Key interface {
	Compare(key Key) int8
	GetNum() int
}
type Node struct {
	Key   Key
	Value interface{}
	Next  *Node
}

func NewNode(Key Key, Value interface{}) *Node {
	return &Node{Key, Value, nil}
}

type Index struct {
	Node  *Node
	Right *Index
	Down  *Index
	Level int
}

func NewIndex(Node *Node, Down *Index, Level int) *Index {
	return &Index{Node, nil, Down, Level}
}

var BasicLevel int = 0
var SupporseRange int = 5

var randomSeed int = rand.Intn(100) | 0x0100

type SkipList struct {
	head   *Index
	length int
}

func NewSkipList() *SkipList {
	return &SkipList{nil, 0}
}

func (this *SkipList) Get(key Key) *Node {
	index := this.head
	for index != nil {
		compare := index.Node.Key.Compare(key)
		if compare == 0 {
			return index.Node
		} else if compare > 0 {
			return nil
		} else {
			now := index
			for now.Right == nil {
				now = now.Down
				if now == nil {
					return nil
				}
			}
			index = now.Right
		}
	}
	return nil
}

// head is nil
// key < head
// key>=head
func (this *SkipList) Insert(key Key, Value interface{}) *Node {
	if this.head == nil {
		level := this.randomLevel()
		node := NewNode(key, Value)
		var index *Index = nil
		for i := 0; i <= level; i++ {
			index = NewIndex(node, index, i)
		}
		this.head = index
		this.length += 1
		return node
	} else {
		if this.head.Node.Key.Compare(key) > 0 {
			level := this.randomLevel()
			if level < this.head.Level {
				level = this.head.Level
			}
			node := NewNode(key, Value)
			node.Next = this.head.Node
			var index *Index = nil
			for i := 0; i <= level; i++ {
				index = NewIndex(node, index, i)
			}
			right := this.head
			this.head = index
			for index != nil {
				if index.Level > right.Level {
					index = index.Down
					continue
				}
				index.Right = right
				right = right.Down
				index = index.Down
			}
			this.length += 1
			return node
		} else {
			oldNode := this.Get(key)
			if oldNode != nil {
				oldNode.Value = Value
				return oldNode
			}
			this.length += 1
			level := this.randomLevel()
			if level > this.head.Level {
				newHead := this.head
				for i := this.head.Level + 1; i <= level; i++ {
					newHead = NewIndex(newHead.Node, newHead, i)
				}
				this.head = newHead
			}
			index := this.head
			newNode := NewNode(key, Value)
			newIndexArray := make([]*Index, 0, level+1)
			var newIndex *Index = nil
			for i := 0; i <= level; i++ {
				newIndex = NewIndex(newNode, newIndex, i)
				newIndexArray = append(newIndexArray, newIndex)
			}
			for index != nil {
				if index.Level > level {
					index = index.Down
					continue
				}
				if index.Right == nil {
					index.Right = newIndexArray[index.Level]
					index = index.Down
				} else {
					compare := index.Right.Node.Key.Compare(key)
					if compare < 0 {
						index = index.Right
					} else {
						newIndex = newIndexArray[index.Level]
						newIndex.Right = index.Right
						index.Right = newIndex
						index = index.Down
					}
				}
			}
			if newIndexArray[0].Right != nil {
				newNode.Next = newIndexArray[0].Right.Node
			}
			return newNode
		}
	}
}

// if not exist
// if head
// other
func (this *SkipList) Delete(key Key) *Node {
	oldNode := this.Get(key)
	if oldNode == nil {
		return nil
	}
	this.length -= 1
	if this.head.Node.Key.Compare(key) == 0 {
		nextNode := this.head.Node.Next
		if nextNode == nil {
			oldNode := this.head.Node
			this.head = nil
			return oldNode
		}
		index := this.head
		for index.Down != nil {
			if index.Node.Key.Compare(nextNode.Key) == 0 {
				break
			}
			index = index.Down
		}
		startLevel := index.Level + 1
		fixLength := this.head.Level - index.Level
		newIndexArray := make([]*Index, 0, fixLength)
		newIndex := index.Right
		for i := startLevel; i <= this.head.Level; i++ {
			newIndex := NewIndex(newIndex.Node, newIndex, i)
			newIndexArray = append(newIndexArray, newIndex)
		}
		replaceOldIndex := this.head
		for i := 0; i < fixLength; i++ {
			newIndex := newIndexArray[replaceOldIndex.Level-startLevel]
			newIndex.Right = replaceOldIndex.Right
			replaceOldIndex = replaceOldIndex.Down
		}
		this.head = newIndexArray[len(newIndexArray)-1]
	} else {
		last := this.head
		for last != nil {
			index := last.Right
			if index == nil {
				last = last.Down
			} else {
				compare := index.Node.Key.Compare(key)
				if compare == 0 {
					last.Right = index.Right
					last = last.Down
				} else if compare < 0 {
					last = index
				} else {
					last = last.Down
				}
			}
		}
	}
	return oldNode
}

func (this *SkipList) randomLevel() int {
	x := randomSeed
	x ^= x << 13
	x ^= x >> 16
	x ^= x << 5
	randomSeed = x
	if (x & 0x80000001) != 0 {
		return 0
	}
	level := 1
	x = x >> 1
	for (x & 1) != 0 {
		level += 1
		x = x >> 1
	}
	return level
}
<|endoftext|>


<commit_before>// Copyright 2016 Circonus, Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Acknowledgement API support - Fetch, Create, Update, Delete*, and Search
// See: https://login.circonus.com/resources/api/calls/acknowledgement
// *  : delete (cancel) by updating with AcknowledgedUntil set to 0

package api

import (
	"encoding/json"
	"fmt"
	"net/url"
	"regexp"
)

// Acknowledgement defines a acknowledgement
type Acknowledgement struct {
	CID               string      `json:"_cid,omitempty"`
	AcknowledgedBy    string      `json:"_acknowledged_by,omitempty"`
	AcknowledgedOn    uint        `json:"_acknowledged_on,omitempty"`
	Active            bool        `json:"_active,omitempty"`
	LastModified      uint        `json:"_last_modified,omitempty"`
	LastModifiedBy    string      `json:"_last_modified_by,omitempty"`
	AcknowledgedUntil interface{} `json:"acknowledged_until,omitempty"`
	AlertCID          string      `json:"alert,omitempty"`
	Notes             string      `json:"notes,omitempty"`
}

const (
	baseAcknowledgementPath = "/acknowledgement"
	acknowledgementCIDRegex = "^" + baseAcknowledgementPath + "/[0-9]+$"
)

// FetchAcknowledgement retrieves a acknowledgement definition
func (a *API) FetchAcknowledgement(cid CIDType) (*Acknowledgement, error) {
	if cid == nil || *cid == "" {
		return nil, fmt.Errorf("Invalid acknowledgement CID [none]")
	}

	acknowledgementCID := string(*cid)

	matched, err := regexp.MatchString(acknowledgementCIDRegex, acknowledgementCID)
	if err != nil {
		return nil, err
	}
	if !matched {
		return nil, fmt.Errorf("Invalid acknowledgement CID [%s]", acknowledgementCID)
	}

	result, err := a.Get(acknowledgementCID)
	if err != nil {
		return nil, err
	}

	acknowledgement := &Acknowledgement{}
	if err := json.Unmarshal(result, acknowledgement); err != nil {
		return nil, err
	}

	return acknowledgement, nil
}

// FetchAcknowledgements retrieves all acknowledgements
func (a *API) FetchAcknowledgements() (*[]Acknowledgement, error) {
	result, err := a.Get(baseAcknowledgementPath)
	if err != nil {
		return nil, err
	}

	var acknowledgements []Acknowledgement
	if err := json.Unmarshal(result, &acknowledgements); err != nil {
		return nil, err
	}

	return &acknowledgements, nil
}

// UpdateAcknowledgement update acknowledgement definition
func (a *API) UpdateAcknowledgement(config *Acknowledgement) (*Acknowledgement, error) {
	if config == nil {
		return nil, fmt.Errorf("Invalid acknowledgement config [nil]")
	}

	acknowledgementCID := string(config.CID)

	matched, err := regexp.MatchString(acknowledgementCIDRegex, acknowledgementCID)
	if err != nil {
		return nil, err
	}
	if !matched {
		return nil, fmt.Errorf("Invalid acknowledgement CID [%s]", acknowledgementCID)
	}

	cfg, err := json.Marshal(config)
	if err != nil {
		return nil, err
	}

	result, err := a.Put(acknowledgementCID, cfg)
	if err != nil {
		return nil, err
	}

	acknowledgement := &Acknowledgement{}
	if err := json.Unmarshal(result, acknowledgement); err != nil {
		return nil, err
	}

	return acknowledgement, nil
}

// CreateAcknowledgement create a new acknowledgement
func (a *API) CreateAcknowledgement(config *Acknowledgement) (*Acknowledgement, error) {
	if config == nil {
		return nil, fmt.Errorf("Invalid acknowledgement config [nil]")
	}

	cfg, err := json.Marshal(config)
	if err != nil {
		return nil, err
	}

	result, err := a.Post(baseAcknowledgementPath, cfg)
	if err != nil {
		return nil, err
	}

	acknowledgement := &Acknowledgement{}
	if err := json.Unmarshal(result, acknowledgement); err != nil {
		return nil, err
	}

	return acknowledgement, nil
}

// SearchAcknowledgements returns list of acknowledgements matching a search query and/or filter
//    - a search query (see: https://login.circonus.com/resources/api#searching)
//    - a filter (see: https://login.circonus.com/resources/api#filtering)
func (a *API) SearchAcknowledgements(searchCriteria *SearchQueryType, filterCriteria *SearchFilterType) (*[]Acknowledgement, error) {
	q := url.Values{}

	if searchCriteria != nil && *searchCriteria != "" {
		q.Set("search", string(*searchCriteria))
	}

	if filterCriteria != nil && len(*filterCriteria) > 0 {
		for filter, criteria := range *filterCriteria {
			for _, val := range criteria {
				q.Add(filter, val)
			}
		}
	}

	if q.Encode() == "" {
		return a.FetchAcknowledgements()
	}

	reqURL := url.URL{
		Path:     baseAcknowledgementPath,
		RawQuery: q.Encode(),
	}

	result, err := a.Get(reqURL.String())
	if err != nil {
		return nil, fmt.Errorf("[ERROR] API call error %+v", err)
	}

	var acknowledgements []Acknowledgement
	if err := json.Unmarshal(result, &acknowledgements); err != nil {
		return nil, err
	}

	return &acknowledgements, nil
}
<commit_msg>upd: add comment on AcknowledgedUntil use<commit_after>// Copyright 2016 Circonus, Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Acknowledgement API support - Fetch, Create, Update, Delete*, and Search
// See: https://login.circonus.com/resources/api/calls/acknowledgement
// *  : delete (cancel) by updating with AcknowledgedUntil set to 0

package api

import (
	"encoding/json"
	"fmt"
	"net/url"
	"regexp"
)

// Acknowledgement defines a acknowledgement
type Acknowledgement struct {
	CID               string      `json:"_cid,omitempty"`
	AcknowledgedBy    string      `json:"_acknowledged_by,omitempty"`
	AcknowledgedOn    uint        `json:"_acknowledged_on,omitempty"`
	Active            bool        `json:"_active,omitempty"`
	LastModified      uint        `json:"_last_modified,omitempty"`
	LastModifiedBy    string      `json:"_last_modified_by,omitempty"`
	AcknowledgedUntil interface{} `json:"acknowledged_until,omitempty"` // always received as uint; can be set using string or uint
	AlertCID          string      `json:"alert,omitempty"`
	Notes             string      `json:"notes,omitempty"`
}

const (
	baseAcknowledgementPath = "/acknowledgement"
	acknowledgementCIDRegex = "^" + baseAcknowledgementPath + "/[0-9]+$"
)

// FetchAcknowledgement retrieves a acknowledgement definition
func (a *API) FetchAcknowledgement(cid CIDType) (*Acknowledgement, error) {
	if cid == nil || *cid == "" {
		return nil, fmt.Errorf("Invalid acknowledgement CID [none]")
	}

	acknowledgementCID := string(*cid)

	matched, err := regexp.MatchString(acknowledgementCIDRegex, acknowledgementCID)
	if err != nil {
		return nil, err
	}
	if !matched {
		return nil, fmt.Errorf("Invalid acknowledgement CID [%s]", acknowledgementCID)
	}

	result, err := a.Get(acknowledgementCID)
	if err != nil {
		return nil, err
	}

	acknowledgement := &Acknowledgement{}
	if err := json.Unmarshal(result, acknowledgement); err != nil {
		return nil, err
	}

	return acknowledgement, nil
}

// FetchAcknowledgements retrieves all acknowledgements
func (a *API) FetchAcknowledgements() (*[]Acknowledgement, error) {
	result, err := a.Get(baseAcknowledgementPath)
	if err != nil {
		return nil, err
	}

	var acknowledgements []Acknowledgement
	if err := json.Unmarshal(result, &acknowledgements); err != nil {
		return nil, err
	}

	return &acknowledgements, nil
}

// UpdateAcknowledgement update acknowledgement definition
func (a *API) UpdateAcknowledgement(config *Acknowledgement) (*Acknowledgement, error) {
	if config == nil {
		return nil, fmt.Errorf("Invalid acknowledgement config [nil]")
	}

	acknowledgementCID := string(config.CID)

	matched, err := regexp.MatchString(acknowledgementCIDRegex, acknowledgementCID)
	if err != nil {
		return nil, err
	}
	if !matched {
		return nil, fmt.Errorf("Invalid acknowledgement CID [%s]", acknowledgementCID)
	}

	cfg, err := json.Marshal(config)
	if err != nil {
		return nil, err
	}

	result, err := a.Put(acknowledgementCID, cfg)
	if err != nil {
		return nil, err
	}

	acknowledgement := &Acknowledgement{}
	if err := json.Unmarshal(result, acknowledgement); err != nil {
		return nil, err
	}

	return acknowledgement, nil
}

// CreateAcknowledgement create a new acknowledgement
func (a *API) CreateAcknowledgement(config *Acknowledgement) (*Acknowledgement, error) {
	if config == nil {
		return nil, fmt.Errorf("Invalid acknowledgement config [nil]")
	}

	cfg, err := json.Marshal(config)
	if err != nil {
		return nil, err
	}

	result, err := a.Post(baseAcknowledgementPath, cfg)
	if err != nil {
		return nil, err
	}

	acknowledgement := &Acknowledgement{}
	if err := json.Unmarshal(result, acknowledgement); err != nil {
		return nil, err
	}

	return acknowledgement, nil
}

// SearchAcknowledgements returns list of acknowledgements matching a search query and/or filter
//    - a search query (see: https://login.circonus.com/resources/api#searching)
//    - a filter (see: https://login.circonus.com/resources/api#filtering)
func (a *API) SearchAcknowledgements(searchCriteria *SearchQueryType, filterCriteria *SearchFilterType) (*[]Acknowledgement, error) {
	q := url.Values{}

	if searchCriteria != nil && *searchCriteria != "" {
		q.Set("search", string(*searchCriteria))
	}

	if filterCriteria != nil && len(*filterCriteria) > 0 {
		for filter, criteria := range *filterCriteria {
			for _, val := range criteria {
				q.Add(filter, val)
			}
		}
	}

	if q.Encode() == "" {
		return a.FetchAcknowledgements()
	}

	reqURL := url.URL{
		Path:     baseAcknowledgementPath,
		RawQuery: q.Encode(),
	}

	result, err := a.Get(reqURL.String())
	if err != nil {
		return nil, fmt.Errorf("[ERROR] API call error %+v", err)
	}

	var acknowledgements []Acknowledgement
	if err := json.Unmarshal(result, &acknowledgements); err != nil {
		return nil, err
	}

	return &acknowledgements, nil
}
<|endoftext|>


<commit_before>package utility

import (
	"encoding/hex"
	"io"
)

func ValidateImage(r io.Reader) bool {
	buf := make([]byte, 8)
	num, err := r.Read(buf)
	if err != nil || num < 8 {
		return false
	}

	gif := convertHexSlice([]string{"47", "49", "46"})
	jpg := convertHexSlice([]string{"FF", "D8", "FF", "E0"})
	png := convertHexSlice([]string{"89", "50", "4E", "47", "D", "A", "1A", "A"})

	return byteSliceSubset(buf, gif) || byteSliceSubset(buf, jpg) || byteSliceSubset(buf, png)
}

//b is a subset of a
func byteSliceSubset(a, b []byte) bool {
	for i, el := range b {
		if el != a[i] {
			return false
		}
	}

	return true
}

func convertHexSlice(s []string) []byte {
	b := []byte{}
	for _, hexStr := range s {
		result, err := hex.DecodeString(hexStr)
		if err == nil {
			b = append(b, result[0])
		}
	}

	return b
}<commit_msg>Correctly validate JPGs<commit_after>package utility

import (
	"encoding/hex"
	"io"
)

func ValidateImage(r io.Reader) bool {
	buf := make([]byte, 8)
	num, err := r.Read(buf)
	if err != nil || num < 8 {
		return false
	}

	gif := convertHexSlice([]string{"47", "49", "46"})
	jpg := convertHexSlice([]string{"FF", "D8"})
	png := convertHexSlice([]string{"89", "50", "4E", "47", "D", "A", "1A", "A"})

	return byteSliceSubset(buf, gif) || byteSliceSubset(buf, jpg) || byteSliceSubset(buf, png)
}

//b is a subset of a
func byteSliceSubset(a, b []byte) bool {
	for i, el := range b {
		if el != a[i] {
			return false
		}
	}

	return true
}

func convertHexSlice(s []string) []byte {
	b := []byte{}
	for _, hexStr := range s {
		result, err := hex.DecodeString(hexStr)
		if err == nil {
			b = append(b, result[0])
		}
	}

	return b
}<|endoftext|>


<commit_before>// Copyright (c) 2012, Jorge Acereda Maciá. All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE file.

// Package ffvp8 provides a wrapper around the VP8 codec in ffmpeg.
package ffvp8

// #cgo LDFLAGS: -lavcodec
//
// #include "libavcodec/avcodec.h"
// extern AVCodec ff_vp8_decoder;
import "C"

import (
	"container/list"
	"image"
	"log"
	"reflect"
	"time"
	"unsafe"
)

type Frame struct {
	*image.YCbCr
	Timecode time.Duration
}

func dupPlane(src []byte, stride, w, h int) []byte {
	dst := make([]byte, w*h)
	for i := 0; i < h; i++ {
		drow := i * w
		srow := i * stride
		copy(dst[drow:drow+w], src[srow:])
	}
	return dst
}

func (f *Frame) dup() *Frame {
	w := f.Rect.Dx()
	h := f.Rect.Dy()
	cw := w / 2
	ch := h / 2
	f.Y = dupPlane(f.Y, f.YStride, w, h)
	f.Cb = dupPlane(f.Cb, f.CStride, cw, ch)
	f.Cr = dupPlane(f.Cr, f.CStride, cw, ch)
	f.YStride = w
	f.CStride = cw
	return f
}

func init() {
	C.avcodec_register(&C.ff_vp8_decoder)
}

type Decoder struct {
	c    *C.AVCodec
	cc   *C.AVCodecContext
	imgs list.List
}

func NewDecoder() *Decoder {
	var d Decoder
	d.c = C.avcodec_find_decoder(C.AV_CODEC_ID_VP8)
	d.cc = C.avcodec_alloc_context3(d.c)
	d.cc.opaque = unsafe.Pointer(&d)
	C.avcodec_open2(d.cc, d.c, nil)
	return &d
}

func mkslice(p *C.uint8_t, sz int) []byte {
	slice := make([]byte, 1)
	hdr := (*reflect.SliceHeader)(unsafe.Pointer(&slice))
	hdr.Cap = sz
	hdr.Len = sz
	hdr.Data = uintptr(unsafe.Pointer(p))
	return slice
}

func (d *Decoder) Decode(data []byte, tc time.Duration) *Frame {
	var pkt C.AVPacket
	var fr C.AVFrame
	var got C.int
	C.avcodec_get_frame_defaults(&fr)
	C.av_init_packet(&pkt)
	pkt.data = (*C.uint8_t)(&data[0])
	pkt.size = C.int(len(data))
	if C.avcodec_decode_video2(d.cc, &fr, &got, &pkt) < 0 {
		log.Panic("Unable to decode")
	}
	if got == 0 {
		return nil
	}
	ys := int(fr.linesize[0])
	cs := int(fr.linesize[1])
	yw := int(d.cc.width)
	yh := int(d.cc.height)
	ysz := ys * yh
	csz := cs * yh / 2
	img := &image.YCbCr{
		Y:              mkslice(fr.data[0], ysz),
		Cb:             mkslice(fr.data[1], csz),
		Cr:             mkslice(fr.data[2], csz),
		SubsampleRatio: image.YCbCrSubsampleRatio420,
		YStride:        ys,
		CStride:        cs,
		Rect:           image.Rect(0, 0, yw, yh),
	}
	return (&Frame{img, tc}).dup()
}
<commit_msg>No need to make the slice<commit_after>// Copyright (c) 2012, Jorge Acereda Maciá. All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE file.

// Package ffvp8 provides a wrapper around the VP8 codec in ffmpeg.
package ffvp8

// #cgo LDFLAGS: -lavcodec
//
// #include "libavcodec/avcodec.h"
// extern AVCodec ff_vp8_decoder;
import "C"

import (
	"container/list"
	"image"
	"log"
	"reflect"
	"time"
	"unsafe"
)

type Frame struct {
	*image.YCbCr
	Timecode time.Duration
}

func dupPlane(src []byte, stride, w, h int) []byte {
	dst := make([]byte, w*h)
	for i := 0; i < h; i++ {
		drow := i * w
		srow := i * stride
		copy(dst[drow:drow+w], src[srow:])
	}
	return dst
}

func (f *Frame) dup() *Frame {
	w := f.Rect.Dx()
	h := f.Rect.Dy()
	cw := w / 2
	ch := h / 2
	f.Y = dupPlane(f.Y, f.YStride, w, h)
	f.Cb = dupPlane(f.Cb, f.CStride, cw, ch)
	f.Cr = dupPlane(f.Cr, f.CStride, cw, ch)
	f.YStride = w
	f.CStride = cw
	return f
}

func init() {
	C.avcodec_register(&C.ff_vp8_decoder)
}

type Decoder struct {
	c    *C.AVCodec
	cc   *C.AVCodecContext
	imgs list.List
}

func NewDecoder() *Decoder {
	var d Decoder
	d.c = C.avcodec_find_decoder(C.AV_CODEC_ID_VP8)
	d.cc = C.avcodec_alloc_context3(d.c)
	d.cc.opaque = unsafe.Pointer(&d)
	C.avcodec_open2(d.cc, d.c, nil)
	return &d
}

func mkslice(p *C.uint8_t, sz int) []byte {
	var slice []byte
	hdr := (*reflect.SliceHeader)(unsafe.Pointer(&slice))
	hdr.Cap = sz
	hdr.Len = sz
	hdr.Data = uintptr(unsafe.Pointer(p))
	return slice
}

func (d *Decoder) Decode(data []byte, tc time.Duration) *Frame {
	var pkt C.AVPacket
	var fr C.AVFrame
	var got C.int
	C.avcodec_get_frame_defaults(&fr)
	C.av_init_packet(&pkt)
	pkt.data = (*C.uint8_t)(&data[0])
	pkt.size = C.int(len(data))
	if C.avcodec_decode_video2(d.cc, &fr, &got, &pkt) < 0 {
		log.Panic("Unable to decode")
	}
	if got == 0 {
		return nil
	}
	ys := int(fr.linesize[0])
	cs := int(fr.linesize[1])
	yw := int(d.cc.width)
	yh := int(d.cc.height)
	ysz := ys * yh
	csz := cs * yh / 2
	img := &image.YCbCr{
		Y:              mkslice(fr.data[0], ysz),
		Cb:             mkslice(fr.data[1], csz),
		Cr:             mkslice(fr.data[2], csz),
		SubsampleRatio: image.YCbCrSubsampleRatio420,
		YStride:        ys,
		CStride:        cs,
		Rect:           image.Rect(0, 0, yw, yh),
	}
	return (&Frame{img, tc}).dup()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
	//"fmt"

	"github.com/garyburd/go-oauth/oauth"
)

type Tweet struct {
	Text      string
	TimeStamp string `json:"timestamp_ms"`
	User      struct {
		ScreenName string `json:"screen_name"`
	}
}

var oauthClient = oauth.Client{}

func TwitterInit(configTwitter TwitterConfig) *oauth.Credentials {
	oauthClient.Credentials.Token = configTwitter.ConsumerKey
	oauthClient.Credentials.Secret = configTwitter.ConsumerSecret
	accessToken := configTwitter.AccessToken
	accessTokenSecret := configTwitter.AccessTokenSecret

	var token *oauth.Credentials
	token = &oauth.Credentials{accessToken, accessTokenSecret}

	return token
}

func TwitterPost(token *oauth.Credentials, twText string) error {
	twUrl := "https://api.twitter.com/1.1/statuses/update.json"
	twStatus := []string{twText, strconv.Itoa(int(time.Now().Unix()))}
	twParam := make(url.Values)
	twParam.Set("status", strings.Join(twStatus, ","))

	oauthClient.SignParam(token, "POST", twUrl, twParam)
	res, err := http.PostForm(twUrl, url.Values(twParam))
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return errors.New(strconv.Itoa(res.StatusCode))
	}

	defer res.Body.Close()
	return nil
}

func TwitterStreaming(twitterTimestamp chan string, token *oauth.Credentials, account string) {
	twUrl := "https://userstream.twitter.com/1.1/user.json"
	twParam := make(url.Values)

	oauthClient.SignParam(token, "GET", twUrl, twParam)
	twUrl = twUrl + "?" + twParam.Encode()
	res, err := http.Get(twUrl)
	if err != nil {
		log.Fatalf("failed to get a tweet\n", err)
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		log.Fatalf("failed to get a tweet\n", res.StatusCode)
	}

	buf := bufio.NewReader(res.Body)
	var last []byte
	for {
		b, _, err := buf.ReadLine()
		last = append(last, b...)
		var tweets [1]Tweet
		err = json.Unmarshal(last, &tweets[0])
		if err != nil {
			continue
		}
		last = []byte{}

		for i := len(tweets) - 1; i >= 0; i-- {
			user := tweets[i].User.ScreenName
			//text := tweets[i].Text
			ts := tweets[i].TimeStamp
			if user == account {
				//fmt.Println(user + ": " + text + ": " + ts)
				twitterTimestamp <- ts
			}
		}
	}
}
<commit_msg>refactoring: twitter<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/garyburd/go-oauth/oauth"
)

type Tweet struct {
	Text      string
	TimeStamp string `json:"timestamp_ms"`
	User      struct {
		ScreenName string `json:"screen_name"`
	}
}

var oauthClient = oauth.Client{}

func TwitterInit(configTwitter TwitterConfig) *oauth.Credentials {
	oauthClient.Credentials.Token = configTwitter.ConsumerKey
	oauthClient.Credentials.Secret = configTwitter.ConsumerSecret
	accessToken := configTwitter.AccessToken
	accessTokenSecret := configTwitter.AccessTokenSecret

	var token *oauth.Credentials
	token = &oauth.Credentials{accessToken, accessTokenSecret}

	return token
}

func TwitterPost(token *oauth.Credentials, twText string) error {
	twUrl := "https://api.twitter.com/1.1/statuses/update.json"
	twStatus := []string{twText, strconv.Itoa(int(time.Now().Unix()))}
	twParam := make(url.Values)
	twParam.Set("status", strings.Join(twStatus, " : "))

	oauthClient.SignParam(token, "POST", twUrl, twParam)
	res, err := http.PostForm(twUrl, url.Values(twParam))
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return errors.New(strconv.Itoa(res.StatusCode))
	}

	defer res.Body.Close()
	return nil
}

func TwitterStreaming(twitterTimestamp chan string, token *oauth.Credentials, account string) {
	twUrl := "https://userstream.twitter.com/1.1/user.json"
	twParam := make(url.Values)

	oauthClient.SignParam(token, "GET", twUrl, twParam)
	twUrl = twUrl + "?" + twParam.Encode()
	res, err := http.Get(twUrl)
	if err != nil {
		log.Fatalf("failed to get a tweet\n", err)
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		log.Fatalf("failed to get a tweet\n", res.StatusCode)
	}

	buf := bufio.NewReader(res.Body)
	var last []byte
	for {
		b, _, err := buf.ReadLine()
		last = append(last, b...)
		var tweets [1]Tweet
		err = json.Unmarshal(last, &tweets[0])
		if err != nil {
			continue
		}
		last = []byte{}

		for i := len(tweets) - 1; i >= 0; i-- {
			user := tweets[i].User.ScreenName
			ts := tweets[i].TimeStamp
			if err != nil {
				continue
			}

			if user == account {
				twitterTimestamp <- ts
			}
		}
	}
}
<|endoftext|>


<commit_before>package notification

import (
	"context"
	"fmt"
	"testing"

	"github.com/keel-hq/keel/types"
)

type fakeSender struct {
	sent *types.EventNotification

	shouldConfigure bool
	shouldError     error
}

func (s *fakeSender) Configure(*Config) (bool, error) {
	return s.shouldConfigure, nil
}

func (s *fakeSender) Send(event types.EventNotification) error {
	s.sent = &event
	fmt.Println("sending event")
	return s.shouldError
}

func TestSend(t *testing.T) {
	sndr := New(context.Background())

	sndr.Configure(&Config{
		Level:    types.LevelDebug,
		Attempts: 1,
	})

	fs := &fakeSender{
		shouldConfigure: true,
		shouldError:     nil,
	}

	RegisterSender("fakeSender", fs)
	defer sndr.UnregisterSender("fakeSender")

	err := sndr.Send(types.EventNotification{
		Level:   types.LevelInfo,
		Type:    types.NotificationPreDeploymentUpdate,
		Message: "foo",
	})

	if err != nil {
		t.Errorf("unexpected error: %s", err)
	}

	if fs.sent.Message != "foo" {
		t.Errorf("unexpected notification message: %s", fs.sent.Message)
	}

	if fs.sent.Level != types.LevelInfo {
		t.Errorf("unexpected level: %s", fs.sent.Level)
	}
}

// test when configured level is higher than the event
func TestSendLevelNotificationA(t *testing.T) {
	sndr := New(context.Background())

	sndr.Configure(&Config{
		Level:    types.LevelInfo,
		Attempts: 1,
	})

	fs := &fakeSender{
		shouldConfigure: true,
		shouldError:     nil,
	}

	RegisterSender("fakeSender", fs)
	defer sndr.UnregisterSender("fakeSender")

	err := sndr.Send(types.EventNotification{
		Level:   types.LevelDebug,
		Type:    types.NotificationPreDeploymentUpdate,
		Message: "foo",
	})

	if err != nil {
		t.Errorf("unexpected error: %s", err)
	}

	if fs.sent != nil {
		t.Errorf("didn't expect to find sent even for this level")
	}
}

// event level is higher than the configured
func TestSendLevelNotificationB(t *testing.T) {
	sndr := New(context.Background())

	sndr.Configure(&Config{
		Level:    types.LevelInfo,
		Attempts: 1,
	})

	fs := &fakeSender{
		shouldConfigure: true,
		shouldError:     nil,
	}

	RegisterSender("fakeSender", fs)
	defer sndr.UnregisterSender("fakeSender")

	err := sndr.Send(types.EventNotification{
		Level:   types.LevelSuccess,
		Type:    types.NotificationPreDeploymentUpdate,
		Message: "foo",
	})

	if err != nil {
		t.Errorf("unexpected error: %s", err)
	}

	if fs.sent.Message != "foo" {
		t.Errorf("unexpected notification message: %s", fs.sent.Message)
	}

	if fs.sent.Level != types.LevelSuccess {
		t.Errorf("unexpected level: %s", fs.sent.Level)
	}
}
<commit_msg>test for lower level notification<commit_after>package notification

import (
	"context"
	"fmt"
	"testing"

	"github.com/keel-hq/keel/types"
)

type fakeSender struct {
	sent *types.EventNotification

	shouldConfigure bool
	shouldError     error
}

func (s *fakeSender) Configure(*Config) (bool, error) {
	return s.shouldConfigure, nil
}

func (s *fakeSender) Send(event types.EventNotification) error {
	s.sent = &event
	fmt.Println("sending event")
	return s.shouldError
}

func TestSend(t *testing.T) {
	sndr := New(context.Background())

	sndr.Configure(&Config{
		Level:    types.LevelDebug,
		Attempts: 1,
	})

	fs := &fakeSender{
		shouldConfigure: true,
		shouldError:     nil,
	}

	RegisterSender("fakeSender", fs)
	defer sndr.UnregisterSender("fakeSender")

	err := sndr.Send(types.EventNotification{
		Level:   types.LevelInfo,
		Type:    types.NotificationPreDeploymentUpdate,
		Message: "foo",
	})

	if err != nil {
		t.Errorf("unexpected error: %s", err)
	}

	if fs.sent.Message != "foo" {
		t.Errorf("unexpected notification message: %s", fs.sent.Message)
	}

	if fs.sent.Level != types.LevelInfo {
		t.Errorf("unexpected level: %s", fs.sent.Level)
	}
}

// test when configured level is higher than the event
func TestSendLevelNotificationA(t *testing.T) {
	sndr := New(context.Background())

	sndr.Configure(&Config{
		Level:    types.LevelInfo,
		Attempts: 1,
	})

	fs := &fakeSender{
		shouldConfigure: true,
		shouldError:     nil,
	}

	RegisterSender("fakeSender", fs)
	defer sndr.UnregisterSender("fakeSender")

	err := sndr.Send(types.EventNotification{
		Level:   types.LevelDebug,
		Type:    types.NotificationPreDeploymentUpdate,
		Message: "foo",
	})

	if err != nil {
		t.Errorf("unexpected error: %s", err)
	}

	if fs.sent != nil {
		t.Errorf("didn't expect to find sent even for this level")
	}
}

// event level is higher than the configured
func TestSendLevelNotificationB(t *testing.T) {
	sndr := New(context.Background())

	sndr.Configure(&Config{
		Level:    types.LevelInfo,
		Attempts: 1,
	})

	fs := &fakeSender{
		shouldConfigure: true,
		shouldError:     nil,
	}

	RegisterSender("fakeSender", fs)
	defer sndr.UnregisterSender("fakeSender")

	err := sndr.Send(types.EventNotification{
		Level:   types.LevelSuccess,
		Type:    types.NotificationPreDeploymentUpdate,
		Message: "foo",
	})

	if err != nil {
		t.Errorf("unexpected error: %s", err)
	}

	if fs.sent.Message != "foo" {
		t.Errorf("unexpected notification message: %s", fs.sent.Message)
	}

	if fs.sent.Level != types.LevelSuccess {
		t.Errorf("unexpected level: %s", fs.sent.Level)
	}
}

func TestSendLevelNotificationC(t *testing.T) {
	sndr := New(context.Background())

	sndr.Configure(&Config{
		Level:    types.LevelDebug,
		Attempts: 1,
	})

	fs := &fakeSender{
		shouldConfigure: true,
		shouldError:     nil,
	}

	RegisterSender("fakeSender", fs)
	defer sndr.UnregisterSender("fakeSender")

	err := sndr.Send(types.EventNotification{
		Level:   types.LevelDebug,
		Type:    types.NotificationPreDeploymentUpdate,
		Message: "foo",
	})

	if err != nil {
		t.Errorf("unexpected error: %s", err)
	}

	if fs.sent.Message != "foo" {
		t.Errorf("unexpected notification message: %s", fs.sent.Message)
	}

	if fs.sent.Level != types.LevelDebug {
		t.Errorf("unexpected level: %s", fs.sent.Level)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/boivie/lovebeat-go/alert"
	"github.com/boivie/lovebeat-go/backend"
	"github.com/boivie/lovebeat-go/config"
	"github.com/boivie/lovebeat-go/dashboard"
	"github.com/boivie/lovebeat-go/httpapi"
	"github.com/boivie/lovebeat-go/service"
	"github.com/boivie/lovebeat-go/tcpapi"
	"github.com/boivie/lovebeat-go/udpapi"
	"github.com/gorilla/mux"
	"github.com/op/go-logging"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
)

var log = logging.MustGetLogger("lovebeat")

const (
	VERSION                 = "0.1.0"
	MAX_UNPROCESSED_PACKETS = 1000
)

var (
	udpAddr     = flag.String("udp", ":8127", "UDP service address")
	tcpAddr     = flag.String("tcp", ":8127", "TCP service address")
	debug       = flag.Bool("debug", false, "print statistics sent to graphite")
	showVersion = flag.Bool("version", false, "print version string")
	workDir     = flag.String("workdir", "work", "working directory")
	cfgFile     = flag.String("config", "/etc/lovebeat.cfg", "configuration file")
)

var (
	signalchan = make(chan os.Signal, 1)
)

func signalHandler(be backend.Backend) {
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			be.Sync()
			return
		}
	}
}

func httpServer(port int16, svcs *service.Services) {
	rtr := mux.NewRouter()
	httpapi.Register(rtr, svcs.GetClient())
	dashboard.Register(rtr, svcs.GetClient())
	http.Handle("/", rtr)
	log.Info("HTTP server running on port %d\n", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port), nil)
}

func getHostname() string {
	var hostname, err = os.Hostname()
	if err != nil {
		return fmt.Sprintf("unknown_%d", os.Getpid())
	}
	return strings.Split(hostname, ".")[0]
}

func main() {
	flag.Parse()

	var format = logging.MustStringFormatter("%{level} %{message}")
	logging.SetFormatter(format)
	if *debug {
		logging.SetLevel(logging.DEBUG, "lovebeat")
	} else {
		logging.SetLevel(logging.INFO, "lovebeat")
	}
	log.Debug("Debug logs enabled")

	if *showVersion {
		fmt.Printf("lovebeats v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}

	var cfg = config.ReadConfig(*cfgFile)

	var hostname = getHostname()
	log.Info("Lovebeat v%s started as host %s, PID %d", VERSION, hostname, os.Getpid())

	var be = backend.NewFileBackend(*workDir)
	var alerters = []alert.Alerter{alert.NewMailAlerter(&cfg.Mail),
		alert.NewWebhooksAlerter()}
	var svcs = service.NewServices(be, alerters)

	signal.Notify(signalchan, syscall.SIGTERM)
	signal.Notify(signalchan, os.Interrupt)

	go svcs.Monitor()
	go httpServer(8080, svcs)
	go udpapi.Listener(*udpAddr, svcs.GetClient())
	go tcpapi.Listener(*tcpAddr, svcs.GetClient())

	// Ensure that the 'all' view exists
	svcs.GetClient().CreateOrUpdateView("all", "", "", "")

	log.Info("Ready to handle incoming connections")

	signalHandler(be)
}
<commit_msg>Updated help text<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/boivie/lovebeat-go/alert"
	"github.com/boivie/lovebeat-go/backend"
	"github.com/boivie/lovebeat-go/config"
	"github.com/boivie/lovebeat-go/dashboard"
	"github.com/boivie/lovebeat-go/httpapi"
	"github.com/boivie/lovebeat-go/service"
	"github.com/boivie/lovebeat-go/tcpapi"
	"github.com/boivie/lovebeat-go/udpapi"
	"github.com/gorilla/mux"
	"github.com/op/go-logging"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
)

var log = logging.MustGetLogger("lovebeat")

const (
	VERSION                 = "0.1.0"
	MAX_UNPROCESSED_PACKETS = 1000
)

var (
	udpAddr     = flag.String("udp", ":8127", "UDP service address")
	tcpAddr     = flag.String("tcp", ":8127", "TCP service address")
	debug       = flag.Bool("debug", false, "Enable debug printouts")
	showVersion = flag.Bool("version", false, "Print version string")
	workDir     = flag.String("workdir", "work", "Working directory")
	cfgFile     = flag.String("config", "/etc/lovebeat.cfg", "Configuration file")
)

var (
	signalchan = make(chan os.Signal, 1)
)

func signalHandler(be backend.Backend) {
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			be.Sync()
			return
		}
	}
}

func httpServer(port int16, svcs *service.Services) {
	rtr := mux.NewRouter()
	httpapi.Register(rtr, svcs.GetClient())
	dashboard.Register(rtr, svcs.GetClient())
	http.Handle("/", rtr)
	log.Info("HTTP server running on port %d\n", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port), nil)
}

func getHostname() string {
	var hostname, err = os.Hostname()
	if err != nil {
		return fmt.Sprintf("unknown_%d", os.Getpid())
	}
	return strings.Split(hostname, ".")[0]
}

func main() {
	flag.Parse()

	var format = logging.MustStringFormatter("%{level} %{message}")
	logging.SetFormatter(format)
	if *debug {
		logging.SetLevel(logging.DEBUG, "lovebeat")
	} else {
		logging.SetLevel(logging.INFO, "lovebeat")
	}
	log.Debug("Debug logs enabled")

	if *showVersion {
		fmt.Printf("lovebeats v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}

	var cfg = config.ReadConfig(*cfgFile)

	var hostname = getHostname()
	log.Info("Lovebeat v%s started as host %s, PID %d", VERSION, hostname, os.Getpid())

	var be = backend.NewFileBackend(*workDir)
	var alerters = []alert.Alerter{alert.NewMailAlerter(&cfg.Mail),
		alert.NewWebhooksAlerter()}
	var svcs = service.NewServices(be, alerters)

	signal.Notify(signalchan, syscall.SIGTERM)
	signal.Notify(signalchan, os.Interrupt)

	go svcs.Monitor()
	go httpServer(8080, svcs)
	go udpapi.Listener(*udpAddr, svcs.GetClient())
	go tcpapi.Listener(*tcpAddr, svcs.GetClient())

	// Ensure that the 'all' view exists
	svcs.GetClient().CreateOrUpdateView("all", "", "", "")

	log.Info("Ready to handle incoming connections")

	signalHandler(be)
}
<|endoftext|>


<commit_before>package minhashlsh

import (
	"math/rand"
	"testing"
)

func randomSignature(size int, seed int64) Signature {
	r := rand.New(rand.NewSource(seed))
	sig := make(Signature, size)
	for i := range sig {
		sig[i] = uint64(r.Int63())
	}
	return sig
}

func Test_HashKeyFunc16(t *testing.T) {
	sig := randomSignature(2, 1)
	f := hashKeyFuncGen(2)
	hashKey := f(sig)
	if len(hashKey) != 2*2 {
		t.Fatal(len(hashKey))
	}
}

func Test_HashKeyFunc64(t *testing.T) {
	sig := randomSignature(2, 1)
	f := hashKeyFuncGen(8)
	hashKey := f(sig)
	if len(hashKey) != 8*2 {
		t.Fatal(len(hashKey))
	}
}

func Test_MinhashLSH(t *testing.T) {
	f := NewMinhashLSH16(8, 0.5)
	sig1 := randomSignature(8, 2)
	sig2 := randomSignature(8, 1)
	sig3 := randomSignature(8, 1)

	f.Add("sig1", sig1)
	f.Add("sig2", sig2)
	f.Add("sig3", sig3)
	f.Index()
	for i := range f.hashTables {
		if len(f.hashTables[i]) != 2 {
			t.Fatal(f.hashTables[i])
		}
	}

	found := 0
	for _, key := range f.Query(sig3) {
		if key == "sig2" || key == "sig3" {
			found++
		}
	}
	if found != 2 {
		t.Fatal("unable to retrieve inserted keys")
	}
}
<commit_msg>update test<commit_after>package minhashlsh

import (
	"math/rand"
	"testing"
)

func randomSignature(size int, seed int64) Signature {
	r := rand.New(rand.NewSource(seed))
	sig := make(Signature, size)
	for i := range sig {
		sig[i] = uint64(r.Int63())
	}
	return sig
}

func Test_HashKeyFunc16(t *testing.T) {
	sig := randomSignature(2, 1)
	f := hashKeyFuncGen(2)
	hashKey := f(sig)
	if len(hashKey) != 2*2 {
		t.Fatal(len(hashKey))
	}
}

func Test_HashKeyFunc64(t *testing.T) {
	sig := randomSignature(2, 1)
	f := hashKeyFuncGen(8)
	hashKey := f(sig)
	if len(hashKey) != 8*2 {
		t.Fatal(len(hashKey))
	}
}

func Test_MinhashLSH(t *testing.T) {
	f := NewMinhashLSH16(256, 0.6)
	// sig1 is different from sig2 and sig3
	// sig2 and sig3 are identical
	sig1 := randomSignature(256, 1)
	sig2 := randomSignature(256, 2)
	sig3 := randomSignature(256, 2)

	f.Add("sig1", sig1)
	f.Add("sig2", sig2)
	f.Add("sig3", sig3)
	f.Index()
	// sig1 should be in its own bucket
	// sig2 and sig3 are in another bucket
	for i := range f.hashTables {
		if len(f.hashTables[i]) != 2 {
			t.Fatal(f.hashTables[i])
		}
	}

	found := 0
	for _, key := range f.Query(sig3) {
		if key == "sig3" || key == "sig2" {
			found++
		}
	}
	if found != 2 {
		t.Fatal("unable to retrieve inserted keys")
	}
}
<|endoftext|>