<commit_before>package sse

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"
)

type state uint8

const (
	stateConnecting state = 0
	stateOpen       state = 1
	stateClosed     state = 2
)

// Client is main struct that handles connecting/streaming/etc...
type Client struct {
	cl       *http.Client
	req      *http.Request
	r        io.ReadCloser
	last     string
	retry    time.Duration
	stopChan chan struct{}
	event    chan Event

	add       chan listener
	listeners map[string]func(Event)

	sync.RWMutex
	readyState state
}

type listener struct {
	name string
	f    func(Event)
}

// Event is SSE event data represenation
type Event struct {
	ID    string
	Event string
	Data  string
}

// Config is a struct used to define and override default parameters:
//   Client - *http.Client, if non provided, http.DefaultClient will be used.
//   URL    - URL of SSE stream to connect to. Must be provided. No default
//            value.
//   Retry  - time.Duration of how long should SSE client wait before trying
//            to reconnect after disconnection. Default is 2 seconds.
type Config struct {
	Client *http.Client
	URL    string
	Retry  time.Duration
}

// New creates a client based on a passed Config.
func New(cfg *Config) (*Client, error) {
	if cfg.URL == "" {
		return nil, errors.New("sse: URL config option MUST be provided")
	}
	req, err := http.NewRequest("GET", cfg.URL, nil)
	if err != nil {
		return nil, fmt.Errorf("sse: could not make request to %s: %v", cfg.URL, err)
	}
	req.Header.Set("Cache-Control", "no-cache")
	req.Header.Set("Accept", "text/event-stream")
	retry := 2 * time.Second
	if cfg.Retry != 0 {
		retry = cfg.Retry
	}
	client := &Client{
		readyState: stateConnecting,
		req:        req,
		retry:      retry,
		event:      make(chan Event),
		add:        make(chan listener),
		stopChan:   make(chan struct{}),
		listeners:  make(map[string]func(Event)),
	}
	client.cl = http.DefaultClient
	if cfg.Client != nil {
		client.cl = cfg.Client
	}
	go client.run()
	return client, nil
}

// Connect connects to given SSE endpoint and starts reading the stream and
// transmitting events.
func (c *Client) Connect() {
	go c.connect()
}

// AddListener adds a listener for a given event type. A listener is simple
// callback function that passes Event struct.
func (c *Client) AddListener(event string, f func(Event)) {
	c.add <- listener{
		name: event,
		f:    f,
	}
}

// Stop stops the client of accepting any more stream requests.
func (c *Client) Stop() {
	c.stop()
}

func (c *Client) stop() {
	close(c.stopChan)
	c.Lock()
	c.readyState = stateClosed
	c.Unlock()
}

func (c *Client) run() {
	for {
		select {
		case l := <-c.add:
			c.listeners[l.name] = l.f
		case event := <-c.event:
			f, ok := c.listeners[event.Event]
			if !ok {
				continue
			}
			if event.ID != "" {
				c.last = event.ID
			}
			f(event)
		case <-c.stopChan:
			return
		}
	}
}

func (c *Client) read() {
	c.RLock()
	state := c.readyState
	c.RUnlock()
	switch state {
	case stateOpen:
	case stateConnecting:
		return
	case stateClosed:
		return
	default:
		return
	}
	c.decode()
}

func (c *Client) decode() {
	defer c.r.Close()
	dec := bufio.NewReader(c.r)
	_, err := dec.Peek(1)
	if err == io.ErrUnexpectedEOF {
		err = io.EOF
	}
	if err != nil {
		c.fireErrorAndRecover(err)
		return
	}
	for {
		event := new(Event)
		event.Event = "message"
		for {
			line, err := dec.ReadString('\n')
			if err != nil {
				c.fireErrorAndRecover(err)
				return
			}
			if line == "\n" {
				break
			}
			line = strings.TrimSuffix(line, "\n")
			if strings.HasPrefix(line, ":") {
				continue
			}
			sections := strings.SplitN(line, ":", 2)
			field, value := sections[0], ""
			if len(sections) == 2 {
				value = strings.TrimPrefix(sections[1], " ")
			}
			switch field {
			case "event":
				event.Event = value
			case "data":
				event.Data += value + "\n"
			case "id":
				event.ID = value
			}
		}
		event.Data = strings.TrimSuffix(event.Data, "\n")
		c.event <- *event
	}
}

func (c *Client) connect() {
	c.req.Header.Set("Last-Event-ID", c.last)
	resp, err := c.cl.Do(c.req)
	if err != nil {
		go c.connect()
		return
	}
	// TODO: check other status codes
	switch resp.StatusCode {
	case http.StatusOK:
		// TODO: check content-type
	case http.StatusNoContent:
		c.stop()
		return
	case http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:
		c.reconnect()
		return
	default:
		c.stop()
		return
	}
	c.r = resp.Body
	c.fireOpen()
}

func (c *Client) reconnect() {
	c.Lock()
	c.readyState = stateConnecting
	c.Unlock()
	time.Sleep(c.retry)
	// TODO: also implement exponential backoff delay
	go c.connect()
}

func (c *Client) fireOpen() {
	c.Lock()
	c.readyState = stateOpen
	c.Unlock()
	go c.read()
	event := new(Event)
	event.Event = "open"
	c.event <- *event
}

func (c *Client) fireErrorAndRecover(err error) {
	event := new(Event)
	event.Event = "error"
	event.Data = err.Error()
	c.event <- *event
	c.reconnect()
}
<commit_msg>Set retry message as milliseconds to further reconnect times<commit_after>package sse

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

type state uint8

const (
	stateConnecting state = 0
	stateOpen       state = 1
	stateClosed     state = 2
)

// Client is main struct that handles connecting/streaming/etc...
type Client struct {
	cl        *http.Client
	req       *http.Request
	r         io.ReadCloser
	last      string
	retry     time.Duration
	retryChan chan time.Duration
	stopChan  chan struct{}
	event     chan Event

	add       chan listener
	listeners map[string]func(Event)

	sync.RWMutex
	readyState state
}

type listener struct {
	name string
	f    func(Event)
}

// Event is SSE event data represenation
type Event struct {
	ID    string
	Event string
	Data  string
}

// Config is a struct used to define and override default parameters:
//   Client - *http.Client, if non provided, http.DefaultClient will be used.
//   URL    - URL of SSE stream to connect to. Must be provided. No default
//            value.
//   Retry  - time.Duration of how long should SSE client wait before trying
//            to reconnect after disconnection. Default is 2 seconds.
type Config struct {
	Client *http.Client
	URL    string
	Retry  time.Duration
}

// New creates a client based on a passed Config.
func New(cfg *Config) (*Client, error) {
	if cfg.URL == "" {
		return nil, errors.New("sse: URL config option MUST be provided")
	}
	req, err := http.NewRequest("GET", cfg.URL, nil)
	if err != nil {
		return nil, fmt.Errorf("sse: could not make request to %s: %v", cfg.URL, err)
	}
	req.Header.Set("Cache-Control", "no-cache")
	req.Header.Set("Accept", "text/event-stream")
	retry := 2 * time.Second
	if cfg.Retry != 0 {
		retry = cfg.Retry
	}
	client := &Client{
		readyState: stateConnecting,
		req:        req,
		retry:      retry,
		event:      make(chan Event),
		add:        make(chan listener),
		stopChan:   make(chan struct{}),
		retryChan:  make(chan time.Duration),
		listeners:  make(map[string]func(Event)),
	}
	client.cl = http.DefaultClient
	if cfg.Client != nil {
		client.cl = cfg.Client
	}
	go client.run()
	return client, nil
}

// Connect connects to given SSE endpoint and starts reading the stream and
// transmitting events.
func (c *Client) Connect() {
	go c.connect()
}

// AddListener adds a listener for a given event type. A listener is simple
// callback function that passes Event struct.
func (c *Client) AddListener(event string, f func(Event)) {
	c.add <- listener{
		name: event,
		f:    f,
	}
}

// Stop stops the client of accepting any more stream requests.
func (c *Client) Stop() {
	c.stop()
}

func (c *Client) stop() {
	close(c.stopChan)
	c.Lock()
	c.readyState = stateClosed
	c.Unlock()
}

func (c *Client) run() {
	for {
		select {
		case l := <-c.add:
			c.listeners[l.name] = l.f
		case event := <-c.event:
			f, ok := c.listeners[event.Event]
			if !ok {
				continue
			}
			if event.ID != "" {
				c.last = event.ID
			}
			f(event)
		case val := <-c.retryChan:
			c.retry = val
		case <-c.stopChan:
			return
		}
	}
}

func (c *Client) read() {
	c.RLock()
	state := c.readyState
	c.RUnlock()
	switch state {
	case stateOpen:
	case stateConnecting:
		return
	case stateClosed:
		return
	default:
		return
	}
	c.decode()
}

func (c *Client) decode() {
	defer c.r.Close()
	dec := bufio.NewReader(c.r)
	_, err := dec.Peek(1)
	if err == io.ErrUnexpectedEOF {
		err = io.EOF
	}
	if err != nil {
		c.fireErrorAndRecover(err)
		return
	}
	for {
		event := new(Event)
		event.Event = "message"
		for {
			line, err := dec.ReadString('\n')
			if err != nil {
				c.fireErrorAndRecover(err)
				return
			}
			if line == "\n" {
				break
			}
			line = strings.TrimSuffix(line, "\n")
			if strings.HasPrefix(line, ":") {
				continue
			}
			sections := strings.SplitN(line, ":", 2)
			field, value := sections[0], ""
			if len(sections) == 2 {
				value = strings.TrimPrefix(sections[1], " ")
			}
			switch field {
			case "event":
				event.Event = value
			case "data":
				event.Data += value + "\n"
			case "id":
				event.ID = value
			case "retry":
				retry, err := strconv.Atoi(value)
				if err != nil {
					break
				}
				c.retryChan <- time.Duration(retry) * time.Millisecond
			}
		}
		event.Data = strings.TrimSuffix(event.Data, "\n")
		c.event <- *event
	}
}

func (c *Client) connect() {
	c.req.Header.Set("Last-Event-ID", c.last)
	resp, err := c.cl.Do(c.req)
	if err != nil {
		go c.connect()
		return
	}
	// TODO: check other status codes
	switch resp.StatusCode {
	case http.StatusOK:
		// TODO: check content-type
	case http.StatusNoContent:
		c.stop()
		return
	case http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:
		c.reconnect()
		return
	default:
		c.stop()
		return
	}
	c.r = resp.Body
	c.fireOpen()
}

func (c *Client) reconnect() {
	c.Lock()
	c.readyState = stateConnecting
	c.Unlock()
	time.Sleep(c.retry)
	// TODO: also implement exponential backoff delay
	go c.connect()
}

func (c *Client) fireOpen() {
	c.Lock()
	c.readyState = stateOpen
	c.Unlock()
	go c.read()
	event := new(Event)
	event.Event = "open"
	c.event <- *event
}

func (c *Client) fireErrorAndRecover(err error) {
	event := new(Event)
	event.Event = "error"
	event.Data = err.Error()
	c.event <- *event
	c.reconnect()
}
<|endoftext|>


<commit_before>package registry

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"testing"

	"github.com/dotcloud/docker/utils"
)

var (
	IMAGE_ID = "42d718c941f5c532ac049bf0b0ab53f0062f09a03afd4aa4a02c098e46032b9d"
	TOKEN    = []string{"fake-token"}
	REPO     = "foo42/bar"
)

func spawnTestRegistry(t *testing.T) *Registry {
	authConfig := &AuthConfig{}
	r, err := NewRegistry(authConfig, utils.NewHTTPRequestFactory(), makeURL("/v1/"), true)
	if err != nil {
		t.Fatal(err)
	}
	return r
}

func TestPingRegistryEndpoint(t *testing.T) {
	regInfo, err := pingRegistryEndpoint(makeURL("/v1/"))
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, regInfo.Standalone, true, "Expected standalone to be true (default)")
}

func TestGetRemoteHistory(t *testing.T) {
	r := spawnTestRegistry(t)
	hist, err := r.GetRemoteHistory(IMAGE_ID, makeURL("/v1/"), TOKEN)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(hist), 2, "Expected 2 images in history")
	assertEqual(t, hist[0], IMAGE_ID, "Expected "+IMAGE_ID+"as first ancestry")
	assertEqual(t, hist[1], "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
		"Unexpected second ancestry")
}

func TestLookupRemoteImage(t *testing.T) {
	r := spawnTestRegistry(t)
	found := r.LookupRemoteImage(IMAGE_ID, makeURL("/v1/"), TOKEN)
	assertEqual(t, found, true, "Expected remote lookup to succeed")
	found = r.LookupRemoteImage("abcdef", makeURL("/v1/"), TOKEN)
	assertEqual(t, found, false, "Expected remote lookup to fail")
}

func TestGetRemoteImageJSON(t *testing.T) {
	r := spawnTestRegistry(t)
	json, size, err := r.GetRemoteImageJSON(IMAGE_ID, makeURL("/v1/"), TOKEN)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, size, 154, "Expected size 154")
	if len(json) <= 0 {
		t.Fatal("Expected non-empty json")
	}

	_, _, err = r.GetRemoteImageJSON("abcdef", makeURL("/v1/"), TOKEN)
	if err == nil {
		t.Fatal("Expected image not found error")
	}
}

func TestGetRemoteImageLayer(t *testing.T) {
	r := spawnTestRegistry(t)
	data, err := r.GetRemoteImageLayer(IMAGE_ID, makeURL("/v1/"), TOKEN, 0)
	if err != nil {
		t.Fatal(err)
	}
	if data == nil {
		t.Fatal("Expected non-nil data result")
	}

	_, err = r.GetRemoteImageLayer("abcdef", makeURL("/v1/"), TOKEN, 0)
	if err == nil {
		t.Fatal("Expected image not found error")
	}
}

func TestGetRemoteTags(t *testing.T) {
	r := spawnTestRegistry(t)
	tags, err := r.GetRemoteTags([]string{makeURL("/v1/")}, REPO, TOKEN)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(tags), 1, "Expected one tag")
	assertEqual(t, tags["latest"], IMAGE_ID, "Expected tag latest to map to "+IMAGE_ID)

	_, err = r.GetRemoteTags([]string{makeURL("/v1/")}, "foo42/baz", TOKEN)
	if err == nil {
		t.Fatal("Expected error when fetching tags for bogus repo")
	}
}

func TestGetRepositoryData(t *testing.T) {
	r := spawnTestRegistry(t)
	parsedUrl, err := url.Parse(makeURL("/v1/"))
	if err != nil {
		t.Fatal(err)
	}
	host := "http://" + parsedUrl.Host + "/v1/"
	data, err := r.GetRepositoryData("foo42/bar")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(data.ImgList), 2, "Expected 2 images in ImgList")
	assertEqual(t, len(data.Endpoints), 2,
		fmt.Sprintf("Expected 2 endpoints in Endpoints, found %d instead", len(data.Endpoints)))
	assertEqual(t, data.Endpoints[0], host,
		fmt.Sprintf("Expected first endpoint to be %s but found %s instead", host, data.Endpoints[0]))
	assertEqual(t, data.Endpoints[1], "http://test.example.com/v1/",
		fmt.Sprintf("Expected first endpoint to be http://test.example.com/v1/ but found %s instead", data.Endpoints[1]))

}

func TestPushImageJSONRegistry(t *testing.T) {
	r := spawnTestRegistry(t)
	imgData := &ImgData{
		ID:       "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
		Checksum: "sha256:1ac330d56e05eef6d438586545ceff7550d3bdcb6b19961f12c5ba714ee1bb37",
	}

	err := r.PushImageJSONRegistry(imgData, []byte{0x42, 0xdf, 0x0}, makeURL("/v1/"), TOKEN)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPushImageLayerRegistry(t *testing.T) {
	r := spawnTestRegistry(t)
	layer := strings.NewReader("")
	_, _, err := r.PushImageLayerRegistry(IMAGE_ID, layer, makeURL("/v1/"), TOKEN, []byte{})
	if err != nil {
		t.Fatal(err)
	}
}

func TestResolveRepositoryName(t *testing.T) {
	_, _, err := ResolveRepositoryName("https://github.com/dotcloud/docker")
	assertEqual(t, err, ErrInvalidRepositoryName, "Expected error invalid repo name")
	ep, repo, err := ResolveRepositoryName("fooo/bar")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, IndexServerAddress(), "Expected endpoint to be index server address")
	assertEqual(t, repo, "fooo/bar", "Expected resolved repo to be foo/bar")

	u := makeURL("")[7:]
	ep, repo, err = ResolveRepositoryName(u + "/private/moonbase")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, u, "Expected endpoint to be "+u)
	assertEqual(t, repo, "private/moonbase", "Expected endpoint to be private/moonbase")

	ep, repo, err = ResolveRepositoryName("ubuntu-12.04-base")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, IndexServerAddress(), "Expected endpoint to be "+IndexServerAddress())
	assertEqual(t, repo, "ubuntu-12.04-base", "Expected endpoint to be ubuntu-12.04-base")
}

func TestPushRegistryTag(t *testing.T) {
	r := spawnTestRegistry(t)
	err := r.PushRegistryTag("foo42/bar", IMAGE_ID, "stable", makeURL("/v1/"), TOKEN)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPushImageJSONIndex(t *testing.T) {
	r := spawnTestRegistry(t)
	imgData := []*ImgData{
		{
			ID:       "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
			Checksum: "sha256:1ac330d56e05eef6d438586545ceff7550d3bdcb6b19961f12c5ba714ee1bb37",
		},
		{
			ID:       "42d718c941f5c532ac049bf0b0ab53f0062f09a03afd4aa4a02c098e46032b9d",
			Checksum: "sha256:bea7bf2e4bacd479344b737328db47b18880d09096e6674165533aa994f5e9f2",
		},
	}
	repoData, err := r.PushImageJSONIndex("foo42/bar", imgData, false, nil)
	if err != nil {
		t.Fatal(err)
	}
	if repoData == nil {
		t.Fatal("Expected RepositoryData object")
	}
	repoData, err = r.PushImageJSONIndex("foo42/bar", imgData, true, []string{r.indexEndpoint})
	if err != nil {
		t.Fatal(err)
	}
	if repoData == nil {
		t.Fatal("Expected RepositoryData object")
	}
}

func TestSearchRepositories(t *testing.T) {
	r := spawnTestRegistry(t)
	results, err := r.SearchRepositories("fakequery")
	if err != nil {
		t.Fatal(err)
	}
	if results == nil {
		t.Fatal("Expected non-nil SearchResults object")
	}
	assertEqual(t, results.NumResults, 1, "Expected 1 search results")
	assertEqual(t, results.Query, "fakequery", "Expected 'fakequery' as query")
	assertEqual(t, results.Results[0].StarCount, 42, "Expected 'fakeimage' a ot hae 42 stars")
}

func TestValidRepositoryName(t *testing.T) {
	if err := validateRepositoryName("docker/docker"); err != nil {
		t.Fatal(err)
	}
	if err := validateRepositoryName("docker/Docker"); err == nil {
		t.Log("Repository name should be invalid")
		t.Fail()
	}
	if err := validateRepositoryName("docker///docker"); err == nil {
		t.Log("Repository name should be invalid")
		t.Fail()
	}
}

func TestTrustedLocation(t *testing.T) {
	for _, url := range []string{"http://example.com", "https://example.com:7777", "http://docker.io", "http://test.docker.io", "https://fakedocker.com"} {
		req, _ := http.NewRequest("GET", url, nil)
		if trustedLocation(req) == true {
			t.Fatalf("'%s' shouldn't be detected as a trusted location", url)
		}
	}

	for _, url := range []string{"https://docker.io", "https://test.docker.io:80"} {
		req, _ := http.NewRequest("GET", url, nil)
		if trustedLocation(req) == false {
			t.Fatalf("'%s' should be detected as a trusted location", url)
		}
	}
}

func TestAddRequiredHeadersToRedirectedRequests(t *testing.T) {
	for _, urls := range [][]string{
		{"http://docker.io", "https://docker.com"},
		{"https://foo.docker.io:7777", "http://bar.docker.com"},
		{"https://foo.docker.io", "https://example.com"},
	} {
		reqFrom, _ := http.NewRequest("GET", urls[0], nil)
		reqFrom.Header.Add("Content-Type", "application/json")
		reqFrom.Header.Add("Authorization", "super_secret")
		reqTo, _ := http.NewRequest("GET", urls[1], nil)

		AddRequiredHeadersToRedirectedRequests(reqTo, []*http.Request{reqFrom})

		if len(reqTo.Header) != 1 {
			t.Fatal("Expected 1 headers, got %d", len(reqTo.Header))
		}

		if reqTo.Header.Get("Content-Type") != "application/json" {
			t.Fatal("'Content-Type' should be 'application/json'")
		}

		if reqTo.Header.Get("Authorization") != "" {
			t.Fatal("'Authorization' should be empty")
		}
	}

	for _, urls := range [][]string{
		{"https://docker.io", "https://docker.com"},
		{"https://foo.docker.io:7777", "https://bar.docker.com"},
	} {
		reqFrom, _ := http.NewRequest("GET", urls[0], nil)
		reqFrom.Header.Add("Content-Type", "application/json")
		reqFrom.Header.Add("Authorization", "super_secret")
		reqTo, _ := http.NewRequest("GET", urls[1], nil)

		AddRequiredHeadersToRedirectedRequests(reqTo, []*http.Request{reqFrom})

		if len(reqTo.Header) != 2 {
			t.Fatal("Expected 2 headers, got %d", len(reqTo.Header))
		}

		if reqTo.Header.Get("Content-Type") != "application/json" {
			t.Fatal("'Content-Type' should be 'application/json'")
		}

		if reqTo.Header.Get("Authorization") != "super_secret" {
			t.Fatal("'Authorization' should be 'super_secret'")
		}
	}
}
<commit_msg>Fix go vet errors<commit_after>package registry

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"testing"

	"github.com/dotcloud/docker/utils"
)

var (
	IMAGE_ID = "42d718c941f5c532ac049bf0b0ab53f0062f09a03afd4aa4a02c098e46032b9d"
	TOKEN    = []string{"fake-token"}
	REPO     = "foo42/bar"
)

func spawnTestRegistry(t *testing.T) *Registry {
	authConfig := &AuthConfig{}
	r, err := NewRegistry(authConfig, utils.NewHTTPRequestFactory(), makeURL("/v1/"), true)
	if err != nil {
		t.Fatal(err)
	}
	return r
}

func TestPingRegistryEndpoint(t *testing.T) {
	regInfo, err := pingRegistryEndpoint(makeURL("/v1/"))
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, regInfo.Standalone, true, "Expected standalone to be true (default)")
}

func TestGetRemoteHistory(t *testing.T) {
	r := spawnTestRegistry(t)
	hist, err := r.GetRemoteHistory(IMAGE_ID, makeURL("/v1/"), TOKEN)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(hist), 2, "Expected 2 images in history")
	assertEqual(t, hist[0], IMAGE_ID, "Expected "+IMAGE_ID+"as first ancestry")
	assertEqual(t, hist[1], "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
		"Unexpected second ancestry")
}

func TestLookupRemoteImage(t *testing.T) {
	r := spawnTestRegistry(t)
	found := r.LookupRemoteImage(IMAGE_ID, makeURL("/v1/"), TOKEN)
	assertEqual(t, found, true, "Expected remote lookup to succeed")
	found = r.LookupRemoteImage("abcdef", makeURL("/v1/"), TOKEN)
	assertEqual(t, found, false, "Expected remote lookup to fail")
}

func TestGetRemoteImageJSON(t *testing.T) {
	r := spawnTestRegistry(t)
	json, size, err := r.GetRemoteImageJSON(IMAGE_ID, makeURL("/v1/"), TOKEN)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, size, 154, "Expected size 154")
	if len(json) <= 0 {
		t.Fatal("Expected non-empty json")
	}

	_, _, err = r.GetRemoteImageJSON("abcdef", makeURL("/v1/"), TOKEN)
	if err == nil {
		t.Fatal("Expected image not found error")
	}
}

func TestGetRemoteImageLayer(t *testing.T) {
	r := spawnTestRegistry(t)
	data, err := r.GetRemoteImageLayer(IMAGE_ID, makeURL("/v1/"), TOKEN, 0)
	if err != nil {
		t.Fatal(err)
	}
	if data == nil {
		t.Fatal("Expected non-nil data result")
	}

	_, err = r.GetRemoteImageLayer("abcdef", makeURL("/v1/"), TOKEN, 0)
	if err == nil {
		t.Fatal("Expected image not found error")
	}
}

func TestGetRemoteTags(t *testing.T) {
	r := spawnTestRegistry(t)
	tags, err := r.GetRemoteTags([]string{makeURL("/v1/")}, REPO, TOKEN)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(tags), 1, "Expected one tag")
	assertEqual(t, tags["latest"], IMAGE_ID, "Expected tag latest to map to "+IMAGE_ID)

	_, err = r.GetRemoteTags([]string{makeURL("/v1/")}, "foo42/baz", TOKEN)
	if err == nil {
		t.Fatal("Expected error when fetching tags for bogus repo")
	}
}

func TestGetRepositoryData(t *testing.T) {
	r := spawnTestRegistry(t)
	parsedUrl, err := url.Parse(makeURL("/v1/"))
	if err != nil {
		t.Fatal(err)
	}
	host := "http://" + parsedUrl.Host + "/v1/"
	data, err := r.GetRepositoryData("foo42/bar")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(data.ImgList), 2, "Expected 2 images in ImgList")
	assertEqual(t, len(data.Endpoints), 2,
		fmt.Sprintf("Expected 2 endpoints in Endpoints, found %d instead", len(data.Endpoints)))
	assertEqual(t, data.Endpoints[0], host,
		fmt.Sprintf("Expected first endpoint to be %s but found %s instead", host, data.Endpoints[0]))
	assertEqual(t, data.Endpoints[1], "http://test.example.com/v1/",
		fmt.Sprintf("Expected first endpoint to be http://test.example.com/v1/ but found %s instead", data.Endpoints[1]))

}

func TestPushImageJSONRegistry(t *testing.T) {
	r := spawnTestRegistry(t)
	imgData := &ImgData{
		ID:       "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
		Checksum: "sha256:1ac330d56e05eef6d438586545ceff7550d3bdcb6b19961f12c5ba714ee1bb37",
	}

	err := r.PushImageJSONRegistry(imgData, []byte{0x42, 0xdf, 0x0}, makeURL("/v1/"), TOKEN)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPushImageLayerRegistry(t *testing.T) {
	r := spawnTestRegistry(t)
	layer := strings.NewReader("")
	_, _, err := r.PushImageLayerRegistry(IMAGE_ID, layer, makeURL("/v1/"), TOKEN, []byte{})
	if err != nil {
		t.Fatal(err)
	}
}

func TestResolveRepositoryName(t *testing.T) {
	_, _, err := ResolveRepositoryName("https://github.com/dotcloud/docker")
	assertEqual(t, err, ErrInvalidRepositoryName, "Expected error invalid repo name")
	ep, repo, err := ResolveRepositoryName("fooo/bar")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, IndexServerAddress(), "Expected endpoint to be index server address")
	assertEqual(t, repo, "fooo/bar", "Expected resolved repo to be foo/bar")

	u := makeURL("")[7:]
	ep, repo, err = ResolveRepositoryName(u + "/private/moonbase")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, u, "Expected endpoint to be "+u)
	assertEqual(t, repo, "private/moonbase", "Expected endpoint to be private/moonbase")

	ep, repo, err = ResolveRepositoryName("ubuntu-12.04-base")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, IndexServerAddress(), "Expected endpoint to be "+IndexServerAddress())
	assertEqual(t, repo, "ubuntu-12.04-base", "Expected endpoint to be ubuntu-12.04-base")
}

func TestPushRegistryTag(t *testing.T) {
	r := spawnTestRegistry(t)
	err := r.PushRegistryTag("foo42/bar", IMAGE_ID, "stable", makeURL("/v1/"), TOKEN)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPushImageJSONIndex(t *testing.T) {
	r := spawnTestRegistry(t)
	imgData := []*ImgData{
		{
			ID:       "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
			Checksum: "sha256:1ac330d56e05eef6d438586545ceff7550d3bdcb6b19961f12c5ba714ee1bb37",
		},
		{
			ID:       "42d718c941f5c532ac049bf0b0ab53f0062f09a03afd4aa4a02c098e46032b9d",
			Checksum: "sha256:bea7bf2e4bacd479344b737328db47b18880d09096e6674165533aa994f5e9f2",
		},
	}
	repoData, err := r.PushImageJSONIndex("foo42/bar", imgData, false, nil)
	if err != nil {
		t.Fatal(err)
	}
	if repoData == nil {
		t.Fatal("Expected RepositoryData object")
	}
	repoData, err = r.PushImageJSONIndex("foo42/bar", imgData, true, []string{r.indexEndpoint})
	if err != nil {
		t.Fatal(err)
	}
	if repoData == nil {
		t.Fatal("Expected RepositoryData object")
	}
}

func TestSearchRepositories(t *testing.T) {
	r := spawnTestRegistry(t)
	results, err := r.SearchRepositories("fakequery")
	if err != nil {
		t.Fatal(err)
	}
	if results == nil {
		t.Fatal("Expected non-nil SearchResults object")
	}
	assertEqual(t, results.NumResults, 1, "Expected 1 search results")
	assertEqual(t, results.Query, "fakequery", "Expected 'fakequery' as query")
	assertEqual(t, results.Results[0].StarCount, 42, "Expected 'fakeimage' a ot hae 42 stars")
}

func TestValidRepositoryName(t *testing.T) {
	if err := validateRepositoryName("docker/docker"); err != nil {
		t.Fatal(err)
	}
	if err := validateRepositoryName("docker/Docker"); err == nil {
		t.Log("Repository name should be invalid")
		t.Fail()
	}
	if err := validateRepositoryName("docker///docker"); err == nil {
		t.Log("Repository name should be invalid")
		t.Fail()
	}
}

func TestTrustedLocation(t *testing.T) {
	for _, url := range []string{"http://example.com", "https://example.com:7777", "http://docker.io", "http://test.docker.io", "https://fakedocker.com"} {
		req, _ := http.NewRequest("GET", url, nil)
		if trustedLocation(req) == true {
			t.Fatalf("'%s' shouldn't be detected as a trusted location", url)
		}
	}

	for _, url := range []string{"https://docker.io", "https://test.docker.io:80"} {
		req, _ := http.NewRequest("GET", url, nil)
		if trustedLocation(req) == false {
			t.Fatalf("'%s' should be detected as a trusted location", url)
		}
	}
}

func TestAddRequiredHeadersToRedirectedRequests(t *testing.T) {
	for _, urls := range [][]string{
		{"http://docker.io", "https://docker.com"},
		{"https://foo.docker.io:7777", "http://bar.docker.com"},
		{"https://foo.docker.io", "https://example.com"},
	} {
		reqFrom, _ := http.NewRequest("GET", urls[0], nil)
		reqFrom.Header.Add("Content-Type", "application/json")
		reqFrom.Header.Add("Authorization", "super_secret")
		reqTo, _ := http.NewRequest("GET", urls[1], nil)

		AddRequiredHeadersToRedirectedRequests(reqTo, []*http.Request{reqFrom})

		if len(reqTo.Header) != 1 {
			t.Fatalf("Expected 1 headers, got %d", len(reqTo.Header))
		}

		if reqTo.Header.Get("Content-Type") != "application/json" {
			t.Fatal("'Content-Type' should be 'application/json'")
		}

		if reqTo.Header.Get("Authorization") != "" {
			t.Fatal("'Authorization' should be empty")
		}
	}

	for _, urls := range [][]string{
		{"https://docker.io", "https://docker.com"},
		{"https://foo.docker.io:7777", "https://bar.docker.com"},
	} {
		reqFrom, _ := http.NewRequest("GET", urls[0], nil)
		reqFrom.Header.Add("Content-Type", "application/json")
		reqFrom.Header.Add("Authorization", "super_secret")
		reqTo, _ := http.NewRequest("GET", urls[1], nil)

		AddRequiredHeadersToRedirectedRequests(reqTo, []*http.Request{reqFrom})

		if len(reqTo.Header) != 2 {
			t.Fatalf("Expected 2 headers, got %d", len(reqTo.Header))
		}

		if reqTo.Header.Get("Content-Type") != "application/json" {
			t.Fatal("'Content-Type' should be 'application/json'")
		}

		if reqTo.Header.Get("Authorization") != "super_secret" {
			t.Fatal("'Authorization' should be 'super_secret'")
		}
	}
}
<|endoftext|>


<commit_before>package TeleGogo

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

	"github.com/wallnutkraken/TeleGogo/Requests"
)

type client struct {
	token      string
	httpClient *http.Client
}

func (c *client) getToken() string {
	return c.token
}

// GetUpdates receives incoming updates using long polling.
func (c *client) GetUpdates(options GetUpdatesOptions) ([]Update, error) {
	get, err := Requests.CreateBotGetWithArgs(c.token, "getUpdates", options.toArgs()...)
	if err != nil {
		return nil, err
	}
	httpResponse, err := c.httpClient.Do(get)
	if err != nil {
		return nil, err
	}

	responseObj := updateResponse{}
	decoder := json.NewDecoder(httpResponse.Body)

	err = decoder.Decode(&responseObj)
	if err != nil {
		return nil, err
	}
	httpResponse.Body.Close()

	return responseObj.Result, err
}

// SetWebhook NOT TESTED. Use this method to specify a url and receive incoming updates via an outgoing
// webhook. Whenever there is an update for the bot, we will send an HTTPS POST request to the specified url,
// containing a JSON-serialized Update. In case of an unsuccessful request, we will give up after a
// reasonable amount of attempts.
func (c *client) SetWebhook(args SetWebhookArgs) error {
	response, err := c.sendFile(args)
	if err != nil {
		return err
	}

	if response.StatusCode != http.StatusOK {
		return responseToError(response)
	}
	response.Body.Close()

	return nil
}

// DownloadFile downloads the specified file
func (c *client) DownloadFile(file File, path string) error {
	get, err := Requests.CreateFileGet(c.getToken(), file.FilePath)
	if err != nil {
		return err
	}
	resp, err := c.httpClient.Do(get)
	if err != nil {
		return err
	}
	physicalFile, err := os.Create(path)
	if err != nil {
		return err
	}
	defer physicalFile.Close()
	defer resp.Body.Close()

	_, err = io.Copy(physicalFile, resp.Body)
	return err
}

// WhoAmI A simple method for testing your bot's auth token. Requires no parameters.
// Returns basic information about the bot in form of a User object.
func (c *client) WhoAmI() (User, error) {
	request, err := Requests.CreateBotGet(c.token, "getMe")

	if err != nil {
		return User{}, err
	}

	response, err := c.httpClient.Do(request)
	if err != nil {
		return User{}, err
	}

	if response.StatusCode != http.StatusOK {
		return User{}, responseToError(response)
	}
	tgResp := userResponse{}
	decoder := json.NewDecoder(response.Body)

	err = decoder.Decode(&tgResp)
	response.Body.Close()

	return tgResp.Result, err
}

// SendMessage sends a message with the specified arguments. On success returns the sent Message.
func (c *client) SendMessage(args SendMessageArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}

	return responseToMessage(response)
}

func (c *client) ForwardMessage(args ForwardMessageArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}

	return responseToMessage(response)
}

func (c *client) sendNewPhoto(args SendPhotoArgs) (Message, error) {
	response, err := c.sendFile(args)
	if err != nil {
		return Message{}, err
	}

	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	/* Read reply */
	return responseToMessage(response)
}

func (c *client) sendExistingPhoto(args SendPhotoArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	return responseToMessage(response)
}

// SendPhoto Use this method to send photos. On success, the sent Message is returned.
func (c *client) SendPhoto(args SendPhotoArgs) (Message, error) {
	/* Decide whether this is a newly uploaded file or an old one. */
	if args.FileID == "" {
		return c.sendNewPhoto(args)
	}
	return c.sendExistingPhoto(args)
}

func (c *client) sendNewAudio(args SendAudioArgs) (Message, error) {
	response, err := c.sendFile(args)
	if err != nil {
		return Message{}, err
	}

	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	/* Read reply */
	return responseToMessage(response)
}

func (c *client) sendExistingAudio(args SendAudioArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	/* Parse reply */
	return responseToMessage(response)
}

func (c *client) SendAudio(args SendAudioArgs) (Message, error) {
	/* Decide if it's a new or existing file, based on user intent */
	if args.AudioFileID != "" {
		return c.sendNewAudio(args)
	}
	return c.sendExistingAudio(args)
}

func (c *client) resendPhoto(args SendPhotoArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	decoder := json.NewDecoder(response.Body)
	msgResponse := messageReply{}
	if err = decoder.Decode(&msgResponse); err != nil {
		return Message{}, err
	}
	response.Body.Close()

	return msgResponse.Result, nil
}

func responseToMessage(response *http.Response) (Message, error) {
	msg := messageReply{}
	decoder := json.NewDecoder(response.Body)
	err := decoder.Decode(msg)
	defer response.Body.Close()
	return msg.Result, err
}

// NewClient Creates a new Client
func NewClient(token string) (Client, error) {
	c := new(client)
	c.token = token
	c.httpClient = &http.Client{}
	return c, nil
}

// Client represents a bot in Telegram.
type Client interface {
	getToken() string
	DownloadFile(File, string) error
	WhoAmI() (User, error)
	GetUpdates(GetUpdatesOptions) ([]Update, error)
	SendMessage(SendMessageArgs) (Message, error)
	ForwardMessage(ForwardMessageArgs) (Message, error)
	SetWebhook(SetWebhookArgs) error
	SendPhoto(SendPhotoArgs) (Message, error)
	SendAudio(SendAudioArgs) (Message, error)
}
<commit_msg>fixed all JSON message reading failing<commit_after>package TeleGogo

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

	"github.com/wallnutkraken/TeleGogo/Requests"
)

type client struct {
	token      string
	httpClient *http.Client
}

func (c *client) getToken() string {
	return c.token
}

// GetUpdates receives incoming updates using long polling.
func (c *client) GetUpdates(options GetUpdatesOptions) ([]Update, error) {
	get, err := Requests.CreateBotGetWithArgs(c.token, "getUpdates", options.toArgs()...)
	if err != nil {
		return nil, err
	}
	httpResponse, err := c.httpClient.Do(get)
	if err != nil {
		return nil, err
	}

	responseObj := updateResponse{}
	decoder := json.NewDecoder(httpResponse.Body)

	err = decoder.Decode(&responseObj)
	if err != nil {
		return nil, err
	}
	httpResponse.Body.Close()

	return responseObj.Result, err
}

// SetWebhook NOT TESTED. Use this method to specify a url and receive incoming updates via an outgoing
// webhook. Whenever there is an update for the bot, we will send an HTTPS POST request to the specified url,
// containing a JSON-serialized Update. In case of an unsuccessful request, we will give up after a
// reasonable amount of attempts.
func (c *client) SetWebhook(args SetWebhookArgs) error {
	response, err := c.sendFile(args)
	if err != nil {
		return err
	}

	if response.StatusCode != http.StatusOK {
		return responseToError(response)
	}
	response.Body.Close()

	return nil
}

// DownloadFile downloads the specified file
func (c *client) DownloadFile(file File, path string) error {
	get, err := Requests.CreateFileGet(c.getToken(), file.FilePath)
	if err != nil {
		return err
	}
	resp, err := c.httpClient.Do(get)
	if err != nil {
		return err
	}
	physicalFile, err := os.Create(path)
	if err != nil {
		return err
	}
	defer physicalFile.Close()
	defer resp.Body.Close()

	_, err = io.Copy(physicalFile, resp.Body)
	return err
}

// WhoAmI A simple method for testing your bot's auth token. Requires no parameters.
// Returns basic information about the bot in form of a User object.
func (c *client) WhoAmI() (User, error) {
	request, err := Requests.CreateBotGet(c.token, "getMe")

	if err != nil {
		return User{}, err
	}

	response, err := c.httpClient.Do(request)
	if err != nil {
		return User{}, err
	}

	if response.StatusCode != http.StatusOK {
		return User{}, responseToError(response)
	}
	tgResp := userResponse{}
	decoder := json.NewDecoder(response.Body)

	err = decoder.Decode(&tgResp)
	response.Body.Close()

	return tgResp.Result, err
}

// SendMessage sends a message with the specified arguments. On success returns the sent Message.
func (c *client) SendMessage(args SendMessageArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}

	return responseToMessage(response)
}

func (c *client) ForwardMessage(args ForwardMessageArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}

	return responseToMessage(response)
}

func (c *client) sendNewPhoto(args SendPhotoArgs) (Message, error) {
	response, err := c.sendFile(args)
	if err != nil {
		return Message{}, err
	}

	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	/* Read reply */
	return responseToMessage(response)
}

func (c *client) sendExistingPhoto(args SendPhotoArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	return responseToMessage(response)
}

// SendPhoto Use this method to send photos. On success, the sent Message is returned.
func (c *client) SendPhoto(args SendPhotoArgs) (Message, error) {
	/* Decide whether this is a newly uploaded file or an old one. */
	if args.FileID == "" {
		return c.sendNewPhoto(args)
	}
	return c.sendExistingPhoto(args)
}

func (c *client) sendNewAudio(args SendAudioArgs) (Message, error) {
	response, err := c.sendFile(args)
	if err != nil {
		return Message{}, err
	}

	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	/* Read reply */
	return responseToMessage(response)
}

func (c *client) sendExistingAudio(args SendAudioArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	/* Parse reply */
	return responseToMessage(response)
}

func (c *client) SendAudio(args SendAudioArgs) (Message, error) {
	/* Decide if it's a new or existing file, based on user intent */
	if args.AudioPath != "" {
		return c.sendNewAudio(args)
	}
	return c.sendExistingAudio(args)
}

func (c *client) resendPhoto(args SendPhotoArgs) (Message, error) {
	response, err := c.sendJSON(args)
	if err != nil {
		return Message{}, err
	}
	if response.StatusCode != http.StatusOK {
		return Message{}, responseToError(response)
	}
	decoder := json.NewDecoder(response.Body)
	msgResponse := messageReply{}
	if err = decoder.Decode(&msgResponse); err != nil {
		return Message{}, err
	}
	response.Body.Close()

	return msgResponse.Result, nil
}

func responseToMessage(response *http.Response) (Message, error) {
	msg := messageReply{}
	decoder := json.NewDecoder(response.Body)
	err := decoder.Decode(&msg)
	defer response.Body.Close()
	return msg.Result, err
}

// NewClient Creates a new Client
func NewClient(token string) (Client, error) {
	c := new(client)
	c.token = token
	c.httpClient = &http.Client{}
	return c, nil
}

// Client represents a bot in Telegram.
type Client interface {
	getToken() string
	DownloadFile(File, string) error
	WhoAmI() (User, error)
	GetUpdates(GetUpdatesOptions) ([]Update, error)
	SendMessage(SendMessageArgs) (Message, error)
	ForwardMessage(ForwardMessageArgs) (Message, error)
	SetWebhook(SetWebhookArgs) error
	SendPhoto(SendPhotoArgs) (Message, error)
	SendAudio(SendAudioArgs) (Message, error)
}
<|endoftext|>


<commit_before>package eventsocket

type Client struct {
	Id string `json:"Id"`
}

type Clients []*Client

func newClient() (client *Client) {
	client = new(Client)

	id := <-uuidBuilder
	client.Id = id.String()

	clients = append(clients, client)

	return
}

var clients = make(Clients, 16)
<commit_msg>client: added connection upgrader<commit_after>package eventsocket

import (
	"errors"
	"fmt"
	"net/http"

	"github.com/gorilla/websocket"
)

type Client struct {
	Id string          `json:"Id"`
	ws *websocket.Conn `json:-`
}

type Clients map[string]*Client

// the main client store
var clients = make(Clients)

// instantiate a new client, set it's id, and store the client
func newClient() (client *Client) {
	client = new(Client)

	id := <-uuidBuilder
	client.Id = id.String()

	clients[client.Id] = client

	return
}

// fetch a client by it's id
func clientById(id string) (*Client, error) {
	if client, ok := clients[id]; ok {
		return client, nil
	}
	return nil, errors.New(fmt.Sprintf("Client id does not exist: %s", id))
}

// upgrade the http connection to become a ws connection
func (client *Client) connectionUpgrade(w http.ResponseWriter, r *http.Request) error {
	// upgrade the connection
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return err
	}

	// store the connection reference
	client.ws = ws
	return nil
}
<|endoftext|>


<commit_before>package haproxy

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"io"
	"net"
	"strings"

	"github.com/gocarina/gocsv"
)

const (
	socketSchema = "unix:///"
	tcpSchema    = "tcp://"
)

type HAProxyClient struct {
	Addr string
	conn net.Conn
}

func (h *HAProxyClient) RunCommand(cmd string) (*bytes.Buffer, error) {
	err := h.dial()
	if err != nil {
		return nil, err
	}

	done := make(chan bool)
	result := bytes.NewBuffer(nil)

	go func() {
		io.Copy(result, h.conn)
		defer func() { done <- true }()
	}()

	go func() {
		h.conn.Write([]byte(cmd + "\n"))
		defer func() { done <- true }()
	}()

	// Wait for both io streams to close
	for i := 0; i < 2; i++ {
		select {
		case <-done:
		}
	}
	h.conn.Close()

	if strings.HasPrefix(result.String(), "Unknown command") {
		return nil, fmt.Errorf("Unknown command: %s", cmd)
	}

	return result, nil
}

func (h *HAProxyClient) Stats() (services Services, err error) {
	res, err := h.RunCommand("show stat")
	if err != nil {
		return services, err
	}

	allStats := []*Stat{}
	reader := csv.NewReader(res)
	reader.TrailingComma = true
	err = gocsv.UnmarshalCSV(reader, &allStats)
	if err != nil {
		return services, fmt.Errorf("error reading csv: %s", err)
	}

	for _, s := range allStats {
		switch s.SvName {
		case "FRONTEND":
			services.Frontends = append(services.Frontends, s)
		case "BACKEND":
			services.Backends = append(services.Backends, s)
		default:
			services.Listeners = append(services.Listeners, s)
		}
	}

	return services, nil
}

func (h *HAProxyClient) dial() (err error) {
	switch h.schema() {
	case "unix":
		h.conn, err = net.Dial("unix", strings.Replace(h.Addr, socketSchema, "", 1))
	case "tcp":
		h.conn, err = net.Dial("tcp", strings.Replace(h.Addr, tcpSchema, "", 1))
	default:
		return fmt.Errorf("unknown schema")
	}
	return err
}

func (h *HAProxyClient) schema() string {
	if strings.HasPrefix(h.Addr, socketSchema) {
		return "socket"
	}
	if strings.HasPrefix(h.Addr, tcpSchema) {
		return "tcp"
	}
	return ""
}
<commit_msg>add doc strings<commit_after>// Package haproxy provides a minimal client for communicating with, and issuing commands to, HAproxy over a network or file socket.
package haproxy

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"io"
	"net"
	"strings"

	"github.com/gocarina/gocsv"
)

const (
	socketSchema = "unix:///"
	tcpSchema    = "tcp://"
)

// HAProxyClient is the main structure of the library.
type HAProxyClient struct {
	Addr string
	conn net.Conn
}

// Entrypoint to the client. Sends an arbitray command string to HAProxy.
func (h *HAProxyClient) RunCommand(cmd string) (*bytes.Buffer, error) {
	err := h.dial()
	if err != nil {
		return nil, err
	}

	done := make(chan bool)
	result := bytes.NewBuffer(nil)

	go func() {
		io.Copy(result, h.conn)
		defer func() { done <- true }()
	}()

	go func() {
		h.conn.Write([]byte(cmd + "\n"))
		defer func() { done <- true }()
	}()

	// Wait for both io streams to close
	for i := 0; i < 2; i++ {
		select {
		case <-done:
		}
	}
	h.conn.Close()

	if strings.HasPrefix(result.String(), "Unknown command") {
		return nil, fmt.Errorf("Unknown command: %s", cmd)
	}

	return result, nil
}

// Equivalent to HAProxy "show stat" command.
func (h *HAProxyClient) Stats() (services Services, err error) {
	res, err := h.RunCommand("show stat")
	if err != nil {
		return services, err
	}

	allStats := []*Stat{}
	reader := csv.NewReader(res)
	reader.TrailingComma = true
	err = gocsv.UnmarshalCSV(reader, &allStats)
	if err != nil {
		return services, fmt.Errorf("error reading csv: %s", err)
	}

	for _, s := range allStats {
		switch s.SvName {
		case "FRONTEND":
			services.Frontends = append(services.Frontends, s)
		case "BACKEND":
			services.Backends = append(services.Backends, s)
		default:
			services.Listeners = append(services.Listeners, s)
		}
	}

	return services, nil
}

func (h *HAProxyClient) dial() (err error) {
	switch h.schema() {
	case "unix":
		h.conn, err = net.Dial("unix", strings.Replace(h.Addr, socketSchema, "", 1))
	case "tcp":
		h.conn, err = net.Dial("tcp", strings.Replace(h.Addr, tcpSchema, "", 1))
	default:
		return fmt.Errorf("unknown schema")
	}
	return err
}

func (h *HAProxyClient) schema() string {
	if strings.HasPrefix(h.Addr, socketSchema) {
		return "socket"
	}
	if strings.HasPrefix(h.Addr, tcpSchema) {
		return "tcp"
	}
	return ""
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math/rand"
	"net"
	"strings"
)

// Client for connection to relaybroker
type Client struct {
	ID           string
	bot          *bot
	incomingConn net.Conn
	fromClient   chan string
	toClient     chan string
	join         chan string // TODO: this should be some kind of priority queue
	test         []string
}

func newClient(conn net.Conn) Client {
	return Client{
		incomingConn: conn,
		fromClient:   make(chan string, 10),
		toClient:     make(chan string, 10),
		join:         make(chan string, 50000),
		test:         make([]string, 0),
	}
}

func (c *Client) init() {
	go c.joinChannels()
	go c.read()
}

func (c *Client) joinChannels() {
	for channel := range c.join {
		c.bot.join <- channel
	}
}

func (c *Client) read() {
	// cha := make(chan string, 5)
	// go c.relaybrokerCommand(cha)
	for msg := range c.toClient {
		c.incomingConn.Write([]byte(msg + "\r\n"))
		//cha <- msg
	}
	//closeChannel(cha)
}

func closeChannel(c chan string) {
	defer func() {
		if r := recover(); r != nil {
			Log.Error(r)
		}
	}()
	close(c)
}

func (c *Client) close() {
	closeChannel(c.join)
	// keep bot running if he wants to reconnect
	if c.bot.ID != "" {
		// dont let the channel fill up and block
		for m := range c.toClient {
			if c.bot.clientConnected {
				bots[c.ID].toClient <- m
				return
			}
			Log.Debug("msg on dc bot")
		}
	}
	if c.bot.clientConnected {
		return
	}
	closeChannel(c.fromClient)
	closeChannel(c.toClient)
	c.bot.close()
	delete(bots, c.ID)
	Log.Debug("CLOSED CLIENT", c.bot.nick)

}

func (c *Client) handleMessage(line string) {
	c.test = append(c.test, line)
	defer func() {
		if r := recover(); r != nil {
			Log.Error(c.test)
			Log.Error(r)
			c.close()
		}
	}()
	spl := strings.SplitN(line, " ", 2)
	msg := spl[1]
	// irc command
	switch spl[0] {
	case "LOGIN": // log into relaybroker with bot id to enable reconnecting, example: LOGIN pajbot2
		if bot, ok := bots[msg]; ok {
			c.ID = msg
			c.bot = bot
			c.bot.client.toClient = c.toClient
			close(c.join)
			c.join = make(chan string, 50000)
			go c.joinChannels()
			c.bot.clientConnected = true
			Log.Debug("old bot reconnected", msg)
			return
		}
		c.bot = newBot(c)
		c.ID = msg
		c.bot.ID = msg
		c.bot.clientConnected = true
		c.bot.Init()
		bots[msg] = c.bot
	case "PASS":
		pass := msg
		if strings.Contains(msg, ";") {
			passwords := strings.Split(msg, ";")
			pass = passwords[1]
			if cfg.BrokerPass != "" {
				if passwords[0] != cfg.BrokerPass {
					c.toClient <- "invalid relaybroker password\r\n"
					c.close()
					Log.Error("invalid relaybroker password")
					return
				}
			}
		}
		if c.bot == nil {
			c.bot = newBot(c)
			c.bot.Init()
		}
		c.bot.pass = pass
	case "NICK":
		c.bot.nick = strings.ToLower(msg) // make sure the nick is lowercase
		// generate random ID
		if c.bot.ID == "" {
			r := rand.Int31n(123456)
			ID := fmt.Sprintf("%d%s%d", 1, c.bot.nick, r)
			bots[ID] = c.bot
			c.ID = ID
		}
	case "JOIN":
		c.join <- msg
	case "USER":
	default:
		go c.bot.handleMessage(spl)

	}
}
<commit_msg>Temporary fix for bot crash<commit_after>package main

import (
	"fmt"
	"math/rand"
	"net"
	"strings"
)

// Client for connection to relaybroker
type Client struct {
	ID           string
	bot          *bot
	incomingConn net.Conn
	fromClient   chan string
	toClient     chan string
	join         chan string // TODO: this should be some kind of priority queue
	test         []string
}

func newClient(conn net.Conn) Client {
	return Client{
		incomingConn: conn,
		fromClient:   make(chan string, 10),
		toClient:     make(chan string, 10),
		join:         make(chan string, 50000),
		test:         make([]string, 0),
	}
}

func (c *Client) init() {
	go c.joinChannels()
	go c.read()
}

func (c *Client) joinChannels() {
	for channel := range c.join {
		c.bot.join <- channel
	}
}

func (c *Client) read() {
	// cha := make(chan string, 5)
	// go c.relaybrokerCommand(cha)
	for msg := range c.toClient {
		c.incomingConn.Write([]byte(msg + "\r\n"))
		//cha <- msg
	}
	//closeChannel(cha)
}

func closeChannel(c chan string) {
	defer func() {
		if r := recover(); r != nil {
			Log.Error(r)
		}
	}()
	close(c)
}

func (c *Client) close() {
	closeChannel(c.join)
	// keep bot running if he wants to reconnect
	if c.bot.ID != "" {
		// dont let the channel fill up and block
		for m := range c.toClient {
			if c.bot.clientConnected {
				bots[c.ID].toClient <- m
				return
			}
			Log.Debug("msg on dc bot")
		}
	}
	if c.bot.clientConnected {
		return
	}
	closeChannel(c.fromClient)
	closeChannel(c.toClient)
	c.bot.close()
	delete(bots, c.ID)
	Log.Debug("CLOSED CLIENT", c.bot.nick)

}

func (c *Client) handleMessage(line string) {
	c.test = append(c.test, line)
	defer func() {
		if r := recover(); r != nil {
			Log.Error(c.test)
			Log.Error(r)
			c.close()
		}
	}()
	spl := strings.SplitN(line, " ", 2)
	msg := spl[1]
	// irc command
	switch spl[0] {
	case "LOGIN": // log into relaybroker with bot id to enable reconnecting, example: LOGIN pajbot2
		if bot, ok := bots[msg]; ok {
			c.ID = msg
			c.bot = bot
			c.bot.client.toClient = c.toClient
			close(c.join)
			c.join = make(chan string, 50000)
			go c.joinChannels()
			c.bot.clientConnected = true
			Log.Debug("old bot reconnected", msg)
			return
		}
		c.bot = newBot(c)
		c.ID = msg
		c.bot.ID = msg
		c.bot.clientConnected = true
		c.bot.Init()
		bots[msg] = c.bot
	case "PASS":
		pass := msg
		if strings.Contains(msg, ";") {
			passwords := strings.Split(msg, ";")
			pass = passwords[1]
			if cfg.BrokerPass != "" {
				if passwords[0] != cfg.BrokerPass {
					c.toClient <- "invalid relaybroker password\r\n"
					c.close()
					Log.Error("invalid relaybroker password")
					return
				}
			}
		}
		if c.bot == nil {
			c.bot = newBot(c)
			c.bot.Init()
		}
		c.bot.pass = pass
	case "NICK":
		c.bot.nick = strings.ToLower(msg) // make sure the nick is lowercase
		// generate random ID
		if c.bot.ID == "" {
			r := rand.Int31n(123456)
			ID := fmt.Sprintf("%d%s%d", 1, c.bot.nick, r)
			bots[ID] = c.bot
			c.ID = ID
		}
	case "JOIN":
		if c.bot == nil {
			c.bot = newBot(c)
			c.bot.Init()
		}
		c.join <- msg
	case "USER":
	default:
		go c.bot.handleMessage(spl)

	}
}
<|endoftext|>


<commit_before>package dotgit

import (
	"fmt"
	"io"
	"sync/atomic"

	"gopkg.in/src-d/go-git.v4/plumbing"
	"gopkg.in/src-d/go-git.v4/plumbing/format/idxfile"
	"gopkg.in/src-d/go-git.v4/plumbing/format/objfile"
	"gopkg.in/src-d/go-git.v4/plumbing/format/packfile"

	"gopkg.in/src-d/go-billy.v3"
)

// PackWriter is a io.Writer that generates the packfile index simultaneously,
// a packfile.Decoder is used with a file reader to read the file being written
// this operation is synchronized with the write operations.
// The packfile is written in a temp file, when Close is called this file
// is renamed/moved (depends on the Filesystem implementation) to the final
// location, if the PackWriter is not used, nothing is written
type PackWriter struct {
	Notify func(h plumbing.Hash, i idxfile.Idxfile)

	fs       billy.Filesystem
	fr, fw   billy.File
	synced   *syncedReader
	checksum plumbing.Hash
	index    idxfile.Idxfile
	result   chan error
}

func newPackWrite(fs billy.Filesystem) (*PackWriter, error) {
	fw, err := fs.TempFile(fs.Join(objectsPath, packPath), "tmp_pack_")
	if err != nil {
		return nil, err
	}

	fr, err := fs.Open(fw.Name())
	if err != nil {
		return nil, err
	}

	writer := &PackWriter{
		fs:     fs,
		fw:     fw,
		fr:     fr,
		synced: newSyncedReader(fw, fr),
		result: make(chan error),
	}

	go writer.buildIndex()
	return writer, nil
}

func (w *PackWriter) buildIndex() {
	s := packfile.NewScanner(w.synced)
	d, err := packfile.NewDecoder(s, nil)
	if err != nil {
		w.result <- err
		return
	}

	checksum, err := d.Decode()
	if err != nil {
		w.result <- err
		return
	}

	w.checksum = checksum
	w.index.PackfileChecksum = checksum
	w.index.Version = idxfile.VersionSupported

	offsets := d.Offsets()
	for h, crc := range d.CRCs() {
		w.index.Add(h, uint64(offsets[h]), crc)
	}

	w.result <- err
}

// waitBuildIndex waits until buildIndex function finishes, this can terminate
// with a packfile.ErrEmptyPackfile, this means that nothing was written so we
// ignore the error
func (w *PackWriter) waitBuildIndex() error {
	err := <-w.result
	if err == packfile.ErrEmptyPackfile {
		return nil
	}

	return err
}

func (w *PackWriter) Write(p []byte) (int, error) {
	return w.synced.Write(p)
}

// Close closes all the file descriptors and save the final packfile, if nothing
// was written, the tempfiles are deleted without writing a packfile.
func (w *PackWriter) Close() error {
	defer func() {
		if w.Notify != nil {
			w.Notify(w.checksum, w.index)
		}

		close(w.result)
	}()

	if err := w.synced.Close(); err != nil {
		return err
	}

	if err := w.waitBuildIndex(); err != nil {
		return err
	}

	if err := w.fr.Close(); err != nil {
		return err
	}

	if err := w.fw.Close(); err != nil {
		return err
	}

	if len(w.index.Entries) == 0 {
		return w.clean()
	}

	return w.save()
}

func (w *PackWriter) clean() error {
	return w.fs.Remove(w.fw.Name())
}

func (w *PackWriter) save() error {
	base := w.fs.Join(objectsPath, packPath, fmt.Sprintf("pack-%s", w.checksum))
	idx, err := w.fs.Create(fmt.Sprintf("%s.idx", base))
	if err != nil {
		return err
	}

	if err := w.encodeIdx(idx); err != nil {
		return err
	}

	if err := idx.Close(); err != nil {
		return err
	}

	return w.fs.Rename(w.fw.Name(), fmt.Sprintf("%s.pack", base))
}

func (w *PackWriter) encodeIdx(writer io.Writer) error {
	e := idxfile.NewEncoder(writer)
	_, err := e.Encode(&w.index)
	return err
}

type syncedReader struct {
	w io.Writer
	r io.ReadSeeker

	blocked, done uint32
	written, read uint64
	news          chan bool
}

func newSyncedReader(w io.Writer, r io.ReadSeeker) *syncedReader {
	return &syncedReader{
		w:    w,
		r:    r,
		news: make(chan bool),
	}
}

func (s *syncedReader) Write(p []byte) (n int, err error) {
	defer func() {
		written := atomic.AddUint64(&s.written, uint64(n))
		read := atomic.LoadUint64(&s.read)
		if written > read {
			s.wake()
		}
	}()

	n, err = s.w.Write(p)
	return
}

func (s *syncedReader) Read(p []byte) (n int, err error) {
	defer func() { atomic.AddUint64(&s.read, uint64(n)) }()

	for {
		s.sleep()
		n, err = s.r.Read(p)
		if err == io.EOF && !s.isDone() && n == 0 {
			continue
		}

		break
	}

	return
}

func (s *syncedReader) isDone() bool {
	return atomic.LoadUint32(&s.done) == 1
}

func (s *syncedReader) isBlocked() bool {
	return atomic.LoadUint32(&s.blocked) == 1
}

func (s *syncedReader) wake() {
	if s.isBlocked() {
		//	fmt.Println("wake")
		atomic.StoreUint32(&s.blocked, 0)
		s.news <- true
	}
}

func (s *syncedReader) sleep() {
	read := atomic.LoadUint64(&s.read)
	written := atomic.LoadUint64(&s.written)
	if read >= written {
		atomic.StoreUint32(&s.blocked, 1)
		//	fmt.Println("sleep", read, written)
		<-s.news
	}

}

func (s *syncedReader) Seek(offset int64, whence int) (int64, error) {
	if whence == io.SeekCurrent {
		return s.r.Seek(offset, whence)
	}

	p, err := s.r.Seek(offset, whence)
	s.read = uint64(p)

	return p, err
}

func (s *syncedReader) Close() error {
	atomic.StoreUint32(&s.done, 1)
	close(s.news)
	return nil
}

type ObjectWriter struct {
	objfile.Writer
	fs billy.Filesystem
	f  billy.File
}

func newObjectWriter(fs billy.Filesystem) (*ObjectWriter, error) {
	f, err := fs.TempFile(fs.Join(objectsPath, packPath), "tmp_obj_")
	if err != nil {
		return nil, err
	}

	return &ObjectWriter{
		Writer: (*objfile.NewWriter(f)),
		fs:     fs,
		f:      f,
	}, nil
}

func (w *ObjectWriter) Close() error {
	if err := w.Writer.Close(); err != nil {
		return err
	}

	if err := w.f.Close(); err != nil {
		return err
	}

	return w.save()
}

func (w *ObjectWriter) save() error {
	hash := w.Hash().String()
	file := w.fs.Join(objectsPath, hash[0:2], hash[2:40])

	return w.fs.Rename(w.f.Name(), file)
}
<commit_msg>fix race on packfile writer<commit_after>package dotgit

import (
	"fmt"
	"io"
	"sync/atomic"

	"gopkg.in/src-d/go-git.v4/plumbing"
	"gopkg.in/src-d/go-git.v4/plumbing/format/idxfile"
	"gopkg.in/src-d/go-git.v4/plumbing/format/objfile"
	"gopkg.in/src-d/go-git.v4/plumbing/format/packfile"

	"gopkg.in/src-d/go-billy.v3"
)

// PackWriter is a io.Writer that generates the packfile index simultaneously,
// a packfile.Decoder is used with a file reader to read the file being written
// this operation is synchronized with the write operations.
// The packfile is written in a temp file, when Close is called this file
// is renamed/moved (depends on the Filesystem implementation) to the final
// location, if the PackWriter is not used, nothing is written
type PackWriter struct {
	Notify func(h plumbing.Hash, i idxfile.Idxfile)

	fs       billy.Filesystem
	fr, fw   billy.File
	synced   *syncedReader
	checksum plumbing.Hash
	index    idxfile.Idxfile
	result   chan error
}

func newPackWrite(fs billy.Filesystem) (*PackWriter, error) {
	fw, err := fs.TempFile(fs.Join(objectsPath, packPath), "tmp_pack_")
	if err != nil {
		return nil, err
	}

	fr, err := fs.Open(fw.Name())
	if err != nil {
		return nil, err
	}

	writer := &PackWriter{
		fs:     fs,
		fw:     fw,
		fr:     fr,
		synced: newSyncedReader(fw, fr),
		result: make(chan error),
	}

	go writer.buildIndex()
	return writer, nil
}

func (w *PackWriter) buildIndex() {
	s := packfile.NewScanner(w.synced)
	d, err := packfile.NewDecoder(s, nil)
	if err != nil {
		w.result <- err
		return
	}

	checksum, err := d.Decode()
	if err != nil {
		w.result <- err
		return
	}

	w.checksum = checksum
	w.index.PackfileChecksum = checksum
	w.index.Version = idxfile.VersionSupported

	offsets := d.Offsets()
	for h, crc := range d.CRCs() {
		w.index.Add(h, uint64(offsets[h]), crc)
	}

	w.result <- err
}

// waitBuildIndex waits until buildIndex function finishes, this can terminate
// with a packfile.ErrEmptyPackfile, this means that nothing was written so we
// ignore the error
func (w *PackWriter) waitBuildIndex() error {
	err := <-w.result
	if err == packfile.ErrEmptyPackfile {
		return nil
	}

	return err
}

func (w *PackWriter) Write(p []byte) (int, error) {
	return w.synced.Write(p)
}

// Close closes all the file descriptors and save the final packfile, if nothing
// was written, the tempfiles are deleted without writing a packfile.
func (w *PackWriter) Close() error {
	defer func() {
		if w.Notify != nil {
			w.Notify(w.checksum, w.index)
		}

		close(w.result)
	}()

	if err := w.synced.Close(); err != nil {
		return err
	}

	if err := w.waitBuildIndex(); err != nil {
		return err
	}

	if err := w.fr.Close(); err != nil {
		return err
	}

	if err := w.fw.Close(); err != nil {
		return err
	}

	if len(w.index.Entries) == 0 {
		return w.clean()
	}

	return w.save()
}

func (w *PackWriter) clean() error {
	return w.fs.Remove(w.fw.Name())
}

func (w *PackWriter) save() error {
	base := w.fs.Join(objectsPath, packPath, fmt.Sprintf("pack-%s", w.checksum))
	idx, err := w.fs.Create(fmt.Sprintf("%s.idx", base))
	if err != nil {
		return err
	}

	if err := w.encodeIdx(idx); err != nil {
		return err
	}

	if err := idx.Close(); err != nil {
		return err
	}

	return w.fs.Rename(w.fw.Name(), fmt.Sprintf("%s.pack", base))
}

func (w *PackWriter) encodeIdx(writer io.Writer) error {
	e := idxfile.NewEncoder(writer)
	_, err := e.Encode(&w.index)
	return err
}

type syncedReader struct {
	w io.Writer
	r io.ReadSeeker

	blocked, done uint32
	written, read uint64
	news          chan bool
}

func newSyncedReader(w io.Writer, r io.ReadSeeker) *syncedReader {
	return &syncedReader{
		w:    w,
		r:    r,
		news: make(chan bool),
	}
}

func (s *syncedReader) Write(p []byte) (n int, err error) {
	defer func() {
		written := atomic.AddUint64(&s.written, uint64(n))
		read := atomic.LoadUint64(&s.read)
		if written > read {
			s.wake()
		}
	}()

	n, err = s.w.Write(p)
	return
}

func (s *syncedReader) Read(p []byte) (n int, err error) {
	defer func() { atomic.AddUint64(&s.read, uint64(n)) }()

	for {
		s.sleep()
		n, err = s.r.Read(p)
		if err == io.EOF && !s.isDone() && n == 0 {
			continue
		}

		break
	}

	return
}

func (s *syncedReader) isDone() bool {
	return atomic.LoadUint32(&s.done) == 1
}

func (s *syncedReader) isBlocked() bool {
	return atomic.LoadUint32(&s.blocked) == 1
}

func (s *syncedReader) wake() {
	if s.isBlocked() {
		//	fmt.Println("wake")
		atomic.StoreUint32(&s.blocked, 0)
		s.news <- true
	}
}

func (s *syncedReader) sleep() {
	read := atomic.LoadUint64(&s.read)
	written := atomic.LoadUint64(&s.written)
	if read >= written {
		atomic.StoreUint32(&s.blocked, 1)
		//	fmt.Println("sleep", read, written)
		<-s.news
	}

}

func (s *syncedReader) Seek(offset int64, whence int) (int64, error) {
	if whence == io.SeekCurrent {
		return s.r.Seek(offset, whence)
	}

	p, err := s.r.Seek(offset, whence)
	atomic.StoreUint64(&s.read, uint64(p))

	return p, err
}

func (s *syncedReader) Close() error {
	atomic.StoreUint32(&s.done, 1)
	close(s.news)
	return nil
}

type ObjectWriter struct {
	objfile.Writer
	fs billy.Filesystem
	f  billy.File
}

func newObjectWriter(fs billy.Filesystem) (*ObjectWriter, error) {
	f, err := fs.TempFile(fs.Join(objectsPath, packPath), "tmp_obj_")
	if err != nil {
		return nil, err
	}

	return &ObjectWriter{
		Writer: (*objfile.NewWriter(f)),
		fs:     fs,
		f:      f,
	}, nil
}

func (w *ObjectWriter) Close() error {
	if err := w.Writer.Close(); err != nil {
		return err
	}

	if err := w.f.Close(); err != nil {
		return err
	}

	return w.save()
}

func (w *ObjectWriter) save() error {
	hash := w.Hash().String()
	file := w.fs.Join(objectsPath, hash[0:2], hash[2:40])

	return w.fs.Rename(w.f.Name(), file)
}
<|endoftext|>


<commit_before>package alignment

import (
	h "github.com/hivdb/nucamino/scorehandler/general"
	a "github.com/hivdb/nucamino/types/amino"
	// c "github.com/hivdb/nucamino/types/codon"
	n "github.com/hivdb/nucamino/types/nucleic"
	"reflect"
	"testing"
)

const (
	MSG_NOT_EQUAL = "Expect %#v but received %#v"
)

var (
	NSEQ = n.ReadString("ACAGTRTTAGTAGGACCTACACCTGCCAACATAATTGGAAGAAATCTGTTGACYCAG")
	ASEQ = a.ReadString("TVLVGPTPVNIIGRNLLTQ")
)

func TestNewAlignment(t *testing.T) {
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	result := NewAlignment(NSEQ, ASEQ, handler)
	expect := &Alignment{
		q: -1000, r: -200,
		nSeq: NSEQ, aSeq: ASEQ, nSeqLen: 57, aSeqLen: 19,
		scoreHandler: handler, nwMatrix: []int{},
		endPosN: 57, endPosA: 19, maxScore: 9100,
		isSimpleAlignment:             true,
		supportPositionalIndel:        false,
		constIndelCodonOpeningScore:   0,
		constIndelCodonExtensionScore: 200,
	}
	result.report = nil
	if !reflect.DeepEqual(expect, result) {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
}

func TestGetMatrixIndex(t *testing.T) {
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	aln := NewAlignment(NSEQ, ASEQ, handler)
	result := aln.getMatrixIndex(INS, 14, 5)
	expect := 1445
	if result != expect {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
	result = aln.getMatrixIndex(DEL, 14, 5)
	expect = 2605
	if result != expect {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
	result = aln.getMatrixIndex(GENERAL, 14, 5)
	expect = 285
	if result != expect {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
}

func TestGetTypedPos(t *testing.T) {
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	aln := NewAlignment(NSEQ, ASEQ, handler)
	st, pn, pa := aln.getTypedPos(1445)
	if st != INS || pn != 14 || pa != 5 {
		t.Errorf(MSG_NOT_EQUAL, []int{int(INS), 14, 5}, []int{int(st), pn, pa})
	}
	st, pn, pa = aln.getTypedPos(2605)
	if st != DEL || pn != 14 || pa != 5 {
		t.Errorf(MSG_NOT_EQUAL, []int{int(DEL), 14, 5}, []int{int(st), pn, pa})
	}
	st, pn, pa = aln.getTypedPos(285)
	if st != GENERAL || pn != 14 || pa != 5 {
		t.Errorf(MSG_NOT_EQUAL, []int{int(GENERAL), 14, 5}, []int{int(st), pn, pa})
	}
}
<commit_msg>Improve test coverage<commit_after>package alignment

import (
	h "github.com/hivdb/nucamino/scorehandler/general"
	a "github.com/hivdb/nucamino/types/amino"
	f "github.com/hivdb/nucamino/types/frameshift"
	m "github.com/hivdb/nucamino/types/mutation"
	n "github.com/hivdb/nucamino/types/nucleic"
	"reflect"
	"testing"
)

const (
	MSG_NOT_EQUAL = "Expect %#v but received %#v"
)

var (
	NSEQ = n.ReadString("ACAGTRTTAGTAGGACCTACACCTGCCAACATAATTGGAAGAAATCTGTTGACYCAG")
	ASEQ = a.ReadString("TVLVGPTPVNIIGRNLLTQ")
)

func TestNewAlignment(t *testing.T) {
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	result := NewAlignment(NSEQ, ASEQ, handler)
	expect := &Alignment{
		q: -1000, r: -200,
		nSeq: NSEQ, aSeq: ASEQ, nSeqLen: 57, aSeqLen: 19,
		scoreHandler: handler, nwMatrix: []int{},
		endPosN: 57, endPosA: 19, maxScore: 9100,
		isSimpleAlignment:             true,
		supportPositionalIndel:        false,
		constIndelCodonOpeningScore:   0,
		constIndelCodonExtensionScore: 200,
	}
	result.report = nil
	if !reflect.DeepEqual(expect, result) {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
}

func TestNewAlignmentWithIns(t *testing.T) {
	NSEQ_INS := n.ReadString("ACAGTRTTAGTAGGACCTTTTACACCTGCCAACATAATTGGAAGAAATCTGTTGACYCAG")
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	result := NewAlignment(NSEQ_INS, ASEQ, handler)
	expect := &Alignment{
		q: -1000, r: -200,
		nSeq: NSEQ_INS, aSeq: ASEQ, nSeqLen: 60, aSeqLen: 19,
		scoreHandler: handler, nwMatrix: []int{},
		endPosN: 60, endPosA: 19, maxScore: 7700,
		isSimpleAlignment:             false,
		supportPositionalIndel:        false,
		constIndelCodonOpeningScore:   0,
		constIndelCodonExtensionScore: 200,
	}
	result.nwMatrix = []int{}
	result.report = nil
	if !reflect.DeepEqual(expect, result) {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
}

func TestNewAlignmentWithInsFs(t *testing.T) {
	NSEQ_INSFS := n.ReadString("ACAGTRTTAGTAGGACCTTTACACCTGCCAACATAATTGGAAGAAATCTGTTGACYCAG")
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	result := NewAlignment(NSEQ_INSFS, ASEQ, handler)
	expect := &Alignment{
		q: -1000, r: -200,
		nSeq: NSEQ_INSFS, aSeq: ASEQ, nSeqLen: 59, aSeqLen: 19,
		scoreHandler: handler, nwMatrix: []int{},
		endPosN: 59, endPosA: 19, maxScore: 7700,
		isSimpleAlignment:             false,
		supportPositionalIndel:        false,
		constIndelCodonOpeningScore:   0,
		constIndelCodonExtensionScore: 200,
	}
	result.nwMatrix = []int{}
	result.report = nil
	if !reflect.DeepEqual(expect, result) {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
}

func TestNewAlignmentWithDelFs(t *testing.T) {
	NSEQ_DELFS := n.ReadString("AAGTRTTAGTAGGACCTACACCTGCCAACATAATTGGAGAAATCTGTTGACYCAG")
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	result := NewAlignment(NSEQ_DELFS, ASEQ, handler)
	expect := &Alignment{
		q: -1000, r: -200,
		nSeq: NSEQ_DELFS, aSeq: ASEQ, nSeqLen: 55, aSeqLen: 19,
		scoreHandler: handler, nwMatrix: []int{},
		endPosN: 55, endPosA: 19, maxScore: 7000,
		isSimpleAlignment:             false,
		supportPositionalIndel:        false,
		constIndelCodonOpeningScore:   0,
		constIndelCodonExtensionScore: 200,
	}
	result.nwMatrix = []int{}
	result.report = nil
	if !reflect.DeepEqual(expect, result) {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
}

func TestNewAlignmentWithDel(t *testing.T) {
	NSEQ_DEL := n.ReadString("ACAGTRTTAGTAGGACCTACACCTAACATAATTGGAAGAAATCTGTTGACYCA")
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	result := NewAlignment(NSEQ_DEL, ASEQ, handler)
	expect := &Alignment{
		q: -1000, r: -200,
		nSeq: NSEQ_DEL, aSeq: ASEQ, nSeqLen: 53, aSeqLen: 19,
		scoreHandler: handler, nwMatrix: []int{},
		endPosN: 53, endPosA: 19, maxScore: 7450,
		isSimpleAlignment:             false,
		supportPositionalIndel:        false,
		constIndelCodonOpeningScore:   0,
		constIndelCodonExtensionScore: 200,
	}
	result.nwMatrix = []int{}
	result.report = nil
	if !reflect.DeepEqual(expect, result) {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
}

func TestGetMatrixIndex(t *testing.T) {
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	aln := NewAlignment(NSEQ, ASEQ, handler)
	result := aln.getMatrixIndex(INS, 14, 5)
	expect := 1445
	if result != expect {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
	result = aln.getMatrixIndex(DEL, 14, 5)
	expect = 2605
	if result != expect {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
	result = aln.getMatrixIndex(GENERAL, 14, 5)
	expect = 285
	if result != expect {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
}

func TestGetTypedPos(t *testing.T) {
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	aln := NewAlignment(NSEQ, ASEQ, handler)
	st, pn, pa := aln.getTypedPos(1445)
	if st != INS || pn != 14 || pa != 5 {
		t.Errorf(MSG_NOT_EQUAL, []int{int(INS), 14, 5}, []int{int(st), pn, pa})
	}
	st, pn, pa = aln.getTypedPos(2605)
	if st != DEL || pn != 14 || pa != 5 {
		t.Errorf(MSG_NOT_EQUAL, []int{int(DEL), 14, 5}, []int{int(st), pn, pa})
	}
	st, pn, pa = aln.getTypedPos(285)
	if st != GENERAL || pn != 14 || pa != 5 {
		t.Errorf(MSG_NOT_EQUAL, []int{int(GENERAL), 14, 5}, []int{int(st), pn, pa})
	}
}

func TestGetReport(t *testing.T) {
	nseq := n.ReadString("ACAGTRTTAGTAGGACCTACACCTAACATAATTGGAAGAAAAAATCTGTTGACYCA")
	handler := h.New(4, 10, 2, 0, 2, nil, false)
	aln := NewAlignment(nseq, ASEQ, handler)
	result := aln.GetReport()
	expect := &AlignmentReport{
		FirstAA: 1,
		FirstNA: 1,
		LastAA:  18,
		LastNA:  54,
		Mutations: []m.Mutation{
			*m.MakeMutation(9, 25, []n.NucleicAcid{}, a.V),
			*m.MakeMutation(14, 37, []n.NucleicAcid{n.A, n.G, n.A, n.A, n.A, n.A}, a.R),
		},
		FrameShifts: []f.FrameShift{},
		AlignedSites: []AlignedSite{
			AlignedSite{1, 1, 3},
			AlignedSite{2, 4, 3},
			AlignedSite{3, 7, 3},
			AlignedSite{4, 10, 3},
			AlignedSite{5, 13, 3},
			AlignedSite{6, 16, 3},
			AlignedSite{7, 19, 3},
			AlignedSite{8, 22, 3},
			AlignedSite{9, 25, 0},
			AlignedSite{10, 25, 3},
			AlignedSite{11, 28, 3},
			AlignedSite{12, 31, 3},
			AlignedSite{13, 34, 3},
			AlignedSite{14, 37, 6},
			AlignedSite{15, 43, 3},
			AlignedSite{16, 46, 3},
			AlignedSite{17, 49, 3},
			AlignedSite{18, 52, 3},
		},
		AminoAcidsLine:    "T  V  L  V  G  P  T  P  V  N  I  I  G  R     N  L  L  T  ",
		ControlLine:       "::::::::::::::::::::::::---:::::::::::::::+++::::::::::::",
		NucleicAcidsLine:  "ACAGTRTTAGTAGGACCTACACCT   AACATAATTGGAAGAAAAAATCTGTTGACY",
		IsSimpleAlignment: false,
	}
	if !reflect.DeepEqual(result, expect) {
		t.Errorf(MSG_NOT_EQUAL, expect, result)
	}
}
<|endoftext|>


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

package dokan

import (
	"sync"
)

/* Keep Go pointers while passing integers to the C heap.
 *
 * These tables should only be used through the functions
 * defined in this file.
 */

var fsTableLock sync.Mutex
var fsTable = make([]fsTableEntry, 0, 2)
var fiTableLock sync.Mutex
var fiTable = map[uint32]File{}
var fiIdx uint32

type fsTableEntry struct {
	fs        FileSystem
	errChan   chan error
	fileCount uint32
}

func fsTableStore(fs FileSystem, ec chan error) uint32 {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()

	for i, c := range fsTable {
		if c.fs == nil {
			fsTable[i] = fsTableEntry{fs: fs, errChan: ec}
			return uint32(i)
		}
	}

	fsTable = append(fsTable, fsTableEntry{fs: fs, errChan: ec})
	return uint32(len(fsTable) - 1)
}

func fsTableFree(slot uint32) {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()
	if int(slot) < len(fsTable) {
		fsTable[slot] = fsTableEntry{}
	}
}

func fsTableGet(slot uint32) FileSystem {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()
	return fsTable[slot].fs
}

func fsTableGetErrChan(slot uint32) chan error {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()
	return fsTable[slot].errChan
}

func fsTableGetFileCount(slot uint32) uint32 {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()
	return fsTable[slot].fileCount
}

func fiTableStoreFile(global uint32, fi File) uint32 {
	fiTableLock.Lock()
	defer fiTableLock.Unlock()
	for {
		// Just use a simple counter (inside the lock)
		// to look for potential free file handles.
		// Overflowing the counter is ok, but skip
		// counter value zero (for better error detection).
		fiIdx++
		if fiIdx == 0 {
			fiIdx++
		}
		_, exist := fiTable[fiIdx]
		if !exist {
			debug("FID alloc", fiIdx, fi)
			fiTable[fiIdx] = fi
			fsTableLock.Lock()
			defer fsTableLock.Unlock()
			fsTable[global].fileCount++
			return fiIdx
		}
	}
}

func fiTableGetFile(file uint32) File {
	fiTableLock.Lock()
	defer fiTableLock.Unlock()
	var fi = fiTable[file]
	debug("FID get", file, fi)
	return fi
}

func fiTableFreeFile(global uint32, file uint32) {
	fiTableLock.Lock()
	defer fiTableLock.Unlock()
	debug("FID free", global, file, "=>", fiTable[file], "# of open files:", len(fiTable)-1)
	delete(fiTable, file)
	fsTableLock.Lock()
	defer fsTableLock.Unlock()
	fsTable[global].fileCount--
}
<commit_msg>dokan: Simplify locking<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package dokan

import (
	"sync"
)

/* Keep Go pointers while passing integers to the C heap.
 *
 * These tables should only be used through the functions
 * defined in this file.
 */

var fsTableLock sync.Mutex
var fsTable = make([]fsTableEntry, 0, 2)
var fiTableLock sync.Mutex
var fiTable = map[uint32]File{}
var fiIdx uint32

type fsTableEntry struct {
	fs        FileSystem
	errChan   chan error
	fileCount uint32
}

func fsTableStore(fs FileSystem, ec chan error) uint32 {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()

	for i, c := range fsTable {
		if c.fs == nil {
			fsTable[i] = fsTableEntry{fs: fs, errChan: ec}
			return uint32(i)
		}
	}

	fsTable = append(fsTable, fsTableEntry{fs: fs, errChan: ec})
	return uint32(len(fsTable) - 1)
}

func fsTableFree(slot uint32) {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()
	if int(slot) < len(fsTable) {
		fsTable[slot] = fsTableEntry{}
	}
}

func fsTableGet(slot uint32) FileSystem {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()
	return fsTable[slot].fs
}

func fsTableGetErrChan(slot uint32) chan error {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()
	return fsTable[slot].errChan
}

func fsTableGetFileCount(slot uint32) uint32 {
	fsTableLock.Lock()
	defer fsTableLock.Unlock()
	return fsTable[slot].fileCount
}

func fiTableStoreFile(global uint32, fi File) uint32 {
	fsTableLock.Lock()
	fsTable[global].fileCount++
	fsTableLock.Unlock()
	fiTableLock.Lock()
	defer fiTableLock.Unlock()
	for {
		// Just use a simple counter (inside the lock)
		// to look for potential free file handles.
		// Overflowing the counter is ok, but skip
		// counter value zero (for better error detection).
		fiIdx++
		if fiIdx == 0 {
			fiIdx++
		}
		_, exist := fiTable[fiIdx]
		if !exist {
			debug("FID alloc", fiIdx, fi)
			fiTable[fiIdx] = fi
			return fiIdx
		}
	}
}

func fiTableGetFile(file uint32) File {
	fiTableLock.Lock()
	defer fiTableLock.Unlock()
	var fi = fiTable[file]
	debug("FID get", file, fi)
	return fi
}

func fiTableFreeFile(global uint32, file uint32) {
	fsTableLock.Lock()
	fsTable[global].fileCount--
	fsTableLock.Unlock()
	fiTableLock.Lock()
	defer fiTableLock.Unlock()
	debug("FID free", global, file, "=>", fiTable[file], "# of open files:", len(fiTable)-1)
	delete(fiTable, file)
}
<|endoftext|>


<commit_before>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filter"
	subproto "github.com/Symantec/Dominator/proto/sub"
	"path"
	"syscall"
	"time"
)

type state struct {
	subFS                   *filesystem.FileSystem
	requiredFS              *filesystem.FileSystem
	requiredInodeToSubInode map[uint64]uint64
	inodesChanged           map[uint64]bool // Required inode number.
	subFilenameToInode      map[string]uint64
}

func (sub *Sub) buildUpdateRequest(request *subproto.UpdateRequest) {
	fmt.Println("buildUpdateRequest()") // TODO(rgooch): Delete debugging.
	var state state
	state.subFS = &sub.fileSystem.FileSystem
	requiredImage := sub.herd.getImage(sub.requiredImage)
	state.requiredFS = requiredImage.FileSystem
	filter := requiredImage.Filter
	request.Triggers = requiredImage.Triggers
	state.requiredInodeToSubInode = make(map[uint64]uint64)
	state.inodesChanged = make(map[uint64]bool)
	var rusageStart, rusageStop syscall.Rusage
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStart)
	compareDirectories(request, &state,
		&state.subFS.DirectoryInode, &state.requiredFS.DirectoryInode,
		"/", filter)
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStop) // HACK
	cpuTime := time.Duration(rusageStop.Utime.Sec)*time.Second +
		time.Duration(rusageStop.Utime.Usec)*time.Microsecond -
		time.Duration(rusageStart.Utime.Sec)*time.Second -
		time.Duration(rusageStart.Utime.Usec)*time.Microsecond
	fmt.Printf("Build update request took: %s user CPU time\n", cpuTime)
}

func compareDirectories(request *subproto.UpdateRequest, state *state,
	subDirectory, requiredDirectory *filesystem.DirectoryInode,
	myPathName string, filter *filter.Filter) {
	// First look for entries that should be deleted.
	if subDirectory != nil {
		for name := range subDirectory.EntriesByName {
			pathname := path.Join(myPathName, name)
			if filter.Match(pathname) {
				continue
			}
			if _, ok := requiredDirectory.EntriesByName[name]; !ok {
				request.PathsToDelete = append(request.PathsToDelete, pathname)
				fmt.Printf("Delete: %s\n", pathname) // HACK
			}
		}
	}
	for name, requiredEntry := range requiredDirectory.EntriesByName {
		pathname := path.Join(myPathName, name)
		if filter.Match(pathname) {
			continue
		}
		var subEntry *filesystem.DirectoryEntry
		if subDirectory != nil {
			if se, ok := subDirectory.EntriesByName[name]; ok {
				subEntry = se
			}
		}
		if subEntry == nil {
			addEntry(request, state, requiredEntry, pathname)
		} else {
			compareEntries(request, state, subEntry, requiredEntry, pathname,
				filter)
		}
		// If a directory: descend (possibly with the directory for the sub).
		requiredInode := requiredEntry.Inode()
		if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
			var subInode *filesystem.DirectoryInode
			if subEntry != nil {
				if si, ok := subEntry.Inode().(*filesystem.DirectoryInode); ok {
					subInode = si
				}
			}
			compareDirectories(request, state, subInode, requiredInode,
				pathname, filter)
		}
	}
}

func addEntry(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		makeDirectory(request, requiredInode, myPathName, true)
	} else {
		addInode(request, state, requiredEntry, myPathName)
	}
}

func compareEntries(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry,
	myPathName string, filter *filter.Filter) {
	subInode := subEntry.Inode()
	requiredInode := requiredEntry.Inode()
	sameType, sameMetadata, sameData := filesystem.CompareInodes(
		subInode, requiredInode, nil)
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		if sameMetadata {
			return
		}
		if sameType {
			makeDirectory(request, requiredInode, myPathName, false)
		} else {
			request.PathsToDelete = append(request.PathsToDelete, myPathName)
			makeDirectory(request, requiredInode, myPathName, true)
			fmt.Printf("Replace non-directory: %s...\n", myPathName) // HACK
		}
		return
	}
	if sameType && sameData && sameMetadata {
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	if sameType && sameData {
		updateMetadata(request, state, requiredEntry, myPathName)
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	request.PathsToDelete = append(request.PathsToDelete, myPathName)
	addInode(request, state, requiredEntry, myPathName)
}

func relink(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	subInum, ok := state.requiredInodeToSubInode[requiredEntry.InodeNumber]
	if !ok {
		state.requiredInodeToSubInode[requiredEntry.InodeNumber] =
			subEntry.InodeNumber
		return
	}
	if subInum == subEntry.InodeNumber {
		return
	}
	makeHardlink(request,
		myPathName, state.subFS.InodeToFilenamesTable[subInum][0])
}

func makeHardlink(request *subproto.UpdateRequest, source, target string) {
	var hardlink subproto.Hardlink
	hardlink.Source = source
	hardlink.Target = target
	request.HardlinksToMake = append(request.HardlinksToMake, hardlink)
	fmt.Printf("Make link: %s => %s\n", source, target) // HACK
}

func updateMetadata(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	if changed := state.inodesChanged[requiredEntry.InodeNumber]; changed {
		return
	}
	var inode subproto.Inode
	inode.Name = myPathName
	inode.GenericInode = requiredEntry.Inode()
	request.InodesToChange = append(request.InodesToChange, inode)
	state.inodesChanged[requiredEntry.InodeNumber] = true
	fmt.Printf("Update metadata: %s\n", myPathName) // HACK
}

func makeDirectory(request *subproto.UpdateRequest,
	requiredInode *filesystem.DirectoryInode, pathName string, create bool) {
	var newdir subproto.Directory
	newdir.Name = pathName
	newdir.Mode = requiredInode.Mode
	newdir.Uid = requiredInode.Uid
	newdir.Gid = requiredInode.Gid
	if create {
		request.DirectoriesToMake = append(request.DirectoriesToMake, newdir)
		fmt.Printf("Add directory: %s...\n", pathName) // HACK
	} else {
		request.DirectoriesToChange = append(request.DirectoriesToMake, newdir)
		fmt.Printf("Change directory: %s...\n", pathName) // HACK
	}
}

func addInode(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	// Try to find a sibling inode.
	names := state.requiredFS.InodeToFilenamesTable[requiredEntry.InodeNumber]
	if len(names) > 1 {
		var sameDataInode filesystem.GenericInode
		var sameDataName string
		for _, name := range names {
			if inum, found := state.getSubInodeFromFilename(name); found {
				subInode := state.subFS.InodeTable[inum]
				_, sameMetadata, sameData := filesystem.CompareInodes(
					subInode, requiredInode, nil)
				if sameMetadata && sameData {
					makeHardlink(request, name, myPathName)
					return
				}
				if sameData {
					sameDataInode = subInode
					sameDataName = name
				}
			}
		}
		if sameDataInode != nil {
			updateMetadata(request, state, requiredEntry, sameDataName)
			makeHardlink(request, sameDataName, myPathName)
			return
		}
	}
	fmt.Printf("Add entry: %s...\n", myPathName) // HACK
	// TODO(rgooch): Add entry.
}

func (state *state) getSubInodeFromFilename(name string) (uint64, bool) {
	if state.subFilenameToInode == nil {
		fmt.Println("Making subFilenameToInode map...") // HACK
		state.subFilenameToInode = make(map[string]uint64)
		for inum, names := range state.subFS.InodeToFilenamesTable {
			for _, n := range names {
				state.subFilenameToInode[n] = inum
			}
		}
		fmt.Println("Made subFilenameToInode map") // HACK
	}
	inum, ok := state.subFilenameToInode[name]
	return inum, ok
}
<commit_msg>buildUpdateRequest(): add new inodes and make hardlinks to them.<commit_after>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filter"
	subproto "github.com/Symantec/Dominator/proto/sub"
	"path"
	"syscall"
	"time"
)

type state struct {
	subFS                   *filesystem.FileSystem
	requiredFS              *filesystem.FileSystem
	requiredInodeToSubInode map[uint64]uint64
	inodesChanged           map[uint64]bool   // Required inode number.
	inodesCreated           map[uint64]string // Required inode number.
	subFilenameToInode      map[string]uint64
}

func (sub *Sub) buildUpdateRequest(request *subproto.UpdateRequest) {
	fmt.Println("buildUpdateRequest()") // TODO(rgooch): Delete debugging.
	var state state
	state.subFS = &sub.fileSystem.FileSystem
	requiredImage := sub.herd.getImage(sub.requiredImage)
	state.requiredFS = requiredImage.FileSystem
	filter := requiredImage.Filter
	request.Triggers = requiredImage.Triggers
	state.requiredInodeToSubInode = make(map[uint64]uint64)
	state.inodesChanged = make(map[uint64]bool)
	state.inodesCreated = make(map[uint64]string)
	var rusageStart, rusageStop syscall.Rusage
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStart)
	compareDirectories(request, &state,
		&state.subFS.DirectoryInode, &state.requiredFS.DirectoryInode,
		"/", filter)
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStop) // HACK
	cpuTime := time.Duration(rusageStop.Utime.Sec)*time.Second +
		time.Duration(rusageStop.Utime.Usec)*time.Microsecond -
		time.Duration(rusageStart.Utime.Sec)*time.Second -
		time.Duration(rusageStart.Utime.Usec)*time.Microsecond
	fmt.Printf("Build update request took: %s user CPU time\n", cpuTime)
}

func compareDirectories(request *subproto.UpdateRequest, state *state,
	subDirectory, requiredDirectory *filesystem.DirectoryInode,
	myPathName string, filter *filter.Filter) {
	// First look for entries that should be deleted.
	if subDirectory != nil {
		for name := range subDirectory.EntriesByName {
			pathname := path.Join(myPathName, name)
			if filter.Match(pathname) {
				continue
			}
			if _, ok := requiredDirectory.EntriesByName[name]; !ok {
				request.PathsToDelete = append(request.PathsToDelete, pathname)
				fmt.Printf("Delete: %s\n", pathname) // HACK
			}
		}
	}
	for name, requiredEntry := range requiredDirectory.EntriesByName {
		pathname := path.Join(myPathName, name)
		if filter.Match(pathname) {
			continue
		}
		var subEntry *filesystem.DirectoryEntry
		if subDirectory != nil {
			if se, ok := subDirectory.EntriesByName[name]; ok {
				subEntry = se
			}
		}
		if subEntry == nil {
			addEntry(request, state, requiredEntry, pathname)
		} else {
			compareEntries(request, state, subEntry, requiredEntry, pathname,
				filter)
		}
		// If a directory: descend (possibly with the directory for the sub).
		requiredInode := requiredEntry.Inode()
		if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
			var subInode *filesystem.DirectoryInode
			if subEntry != nil {
				if si, ok := subEntry.Inode().(*filesystem.DirectoryInode); ok {
					subInode = si
				}
			}
			compareDirectories(request, state, subInode, requiredInode,
				pathname, filter)
		}
	}
}

func addEntry(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		makeDirectory(request, requiredInode, myPathName, true)
	} else {
		addInode(request, state, requiredEntry, myPathName)
	}
}

func compareEntries(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry,
	myPathName string, filter *filter.Filter) {
	subInode := subEntry.Inode()
	requiredInode := requiredEntry.Inode()
	sameType, sameMetadata, sameData := filesystem.CompareInodes(
		subInode, requiredInode, nil)
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		if sameMetadata {
			return
		}
		if sameType {
			makeDirectory(request, requiredInode, myPathName, false)
		} else {
			request.PathsToDelete = append(request.PathsToDelete, myPathName)
			makeDirectory(request, requiredInode, myPathName, true)
			fmt.Printf("Replace non-directory: %s...\n", myPathName) // HACK
		}
		return
	}
	if sameType && sameData && sameMetadata {
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	if sameType && sameData {
		updateMetadata(request, state, requiredEntry, myPathName)
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	request.PathsToDelete = append(request.PathsToDelete, myPathName)
	addInode(request, state, requiredEntry, myPathName)
}

func relink(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	subInum, ok := state.requiredInodeToSubInode[requiredEntry.InodeNumber]
	if !ok {
		state.requiredInodeToSubInode[requiredEntry.InodeNumber] =
			subEntry.InodeNumber
		return
	}
	if subInum == subEntry.InodeNumber {
		return
	}
	makeHardlink(request,
		myPathName, state.subFS.InodeToFilenamesTable[subInum][0])
}

func makeHardlink(request *subproto.UpdateRequest, source, target string) {
	var hardlink subproto.Hardlink
	hardlink.Source = source
	hardlink.Target = target
	request.HardlinksToMake = append(request.HardlinksToMake, hardlink)
	fmt.Printf("Make link: %s => %s\n", source, target) // HACK
}

func updateMetadata(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	if state.inodesChanged[requiredEntry.InodeNumber] {
		return
	}
	var inode subproto.Inode
	inode.Name = myPathName
	inode.GenericInode = requiredEntry.Inode()
	request.InodesToChange = append(request.InodesToChange, inode)
	state.inodesChanged[requiredEntry.InodeNumber] = true
	fmt.Printf("Update metadata: %s\n", myPathName) // HACK
}

func makeDirectory(request *subproto.UpdateRequest,
	requiredInode *filesystem.DirectoryInode, pathName string, create bool) {
	var newdir subproto.Directory
	newdir.Name = pathName
	newdir.Mode = requiredInode.Mode
	newdir.Uid = requiredInode.Uid
	newdir.Gid = requiredInode.Gid
	if create {
		request.DirectoriesToMake = append(request.DirectoriesToMake, newdir)
		fmt.Printf("Add directory: %s...\n", pathName) // HACK
	} else {
		request.DirectoriesToChange = append(request.DirectoriesToMake, newdir)
		fmt.Printf("Change directory: %s...\n", pathName) // HACK
	}
}

func addInode(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	if name, ok := state.inodesCreated[requiredEntry.InodeNumber]; ok {
		makeHardlink(request, myPathName, name)
		return
	}
	// Try to find a sibling inode.
	names := state.requiredFS.InodeToFilenamesTable[requiredEntry.InodeNumber]
	if len(names) > 1 {
		var sameDataInode filesystem.GenericInode
		var sameDataName string
		for _, name := range names {
			if inum, found := state.getSubInodeFromFilename(name); found {
				subInode := state.subFS.InodeTable[inum]
				_, sameMetadata, sameData := filesystem.CompareInodes(
					subInode, requiredInode, nil)
				if sameMetadata && sameData {
					makeHardlink(request, myPathName, name)
					return
				}
				if sameData {
					sameDataInode = subInode
					sameDataName = name
				}
			}
		}
		if sameDataInode != nil {
			updateMetadata(request, state, requiredEntry, sameDataName)
			makeHardlink(request, myPathName, sameDataName)
			return
		}
	}
	var inode subproto.Inode
	inode.Name = myPathName
	inode.GenericInode = requiredEntry.Inode()
	request.InodesToMake = append(request.InodesToMake, inode)
	state.inodesCreated[requiredEntry.InodeNumber] = myPathName
	fmt.Printf("Add entry: %s...\n", myPathName) // HACK
	// TODO(rgooch): Add entry.
}

func (state *state) getSubInodeFromFilename(name string) (uint64, bool) {
	if state.subFilenameToInode == nil {
		fmt.Println("Making subFilenameToInode map...") // HACK
		state.subFilenameToInode = make(map[string]uint64)
		for inum, names := range state.subFS.InodeToFilenamesTable {
			for _, n := range names {
				state.subFilenameToInode[n] = inum
			}
		}
		fmt.Println("Made subFilenameToInode map") // HACK
	}
	inum, ok := state.subFilenameToInode[name]
	return inum, ok
}
<|endoftext|>


<commit_before>package query

import (
	"fmt"
	"github.com/miekg/dns"
	"math/rand"
	"time"
)

type Result struct {
	Domain string
	Rcode  int
	err    error
}

// Format Result into string for output file
func (dr Result) String(simple bool) string {
	if simple {
		return fmt.Sprintf("%s\n", dr.Domain)
	}
	return fmt.Sprintf("%s\t%s\t%q\t\n", dr.Domain, dns.Rcode_str[dr.Rcode], dr.err)
}

// Return true if the domain is available (DNS NXDOMAIN)
func (dr Result) Available() bool {
	return dr.Rcode == dns.RcodeNameError
}

// Returns true if domain has a Name Server associated
func queryNS(domain string, dnsServers []string, proto string) (int, error) {
	c := new(dns.Client)
	c.ReadTimeout = time.Duration(2 * time.Second)
	c.WriteTimeout = time.Duration(2 * time.Second)
	c.Net = proto
	m := new(dns.Msg)
	m.RecursionDesired = true
	dnsServer := dnsServers[rand.Intn(len(dnsServers))]
	m.SetQuestion(dns.Fqdn(domain), dns.TypeNS)
	in, err := c.Exchange(m, dnsServer+":53")
	if err == nil {
		return in.Rcode, err
	}
	return dns.RcodeRefused, err
}

// Check if each domain
func CheckDomains(id int, in, retries chan string, out chan Result, dnsServers []string, proto string) {
	for {
		var domain string
		select {
		case domain = <-in:
		case domain = <-retries:
		}
		rCode, err := queryNS(domain, dnsServers, proto)
		if err != nil {
			retries <- domain
		} else {
			out <- Result{domain, rCode, err}
		}
	}
}
<commit_msg>Updated code to match miekg DNS pkg<commit_after>package query

import (
	"fmt"
	"github.com/miekg/dns"
	"math/rand"
	"time"
)

type Result struct {
	Domain string
	Rcode  int
	err    error
}

// Format Result into string for output file
func (dr Result) String(simple bool) string {
	if simple {
		return fmt.Sprintf("%s\n", dr.Domain)
	}
	return fmt.Sprintf("%s\t%s\t%q\t\n", dr.Domain, dns.RcodeToString[dr.Rcode], dr.err)
}

// Return true if the domain is available (DNS NXDOMAIN)
func (dr Result) Available() bool {
	return dr.Rcode == dns.RcodeNameError
}

// Returns true if domain has a Name Server associated
func queryNS(domain string, dnsServers []string, proto string) (int, error) {
	c := new(dns.Client)
	c.ReadTimeout = time.Duration(2 * time.Second)
	c.WriteTimeout = time.Duration(2 * time.Second)
	c.Net = proto
	m := new(dns.Msg)
	m.RecursionDesired = true
	dnsServer := dnsServers[rand.Intn(len(dnsServers))]
	m.SetQuestion(dns.Fqdn(domain), dns.TypeNS)
	in, _, err := c.Exchange(m, dnsServer+":53")
	if err == nil {
		return in.Rcode, err
	}
	return dns.RcodeRefused, err
}

// Check if each domain
func CheckDomains(id int, in, retries chan string, out chan Result, dnsServers []string, proto string) {
	for {
		var domain string
		select {
		case domain = <-in:
		case domain = <-retries:
		}
		rCode, err := queryNS(domain, dnsServers, proto)
		if err != nil {
			retries <- domain
		} else {
			out <- Result{domain, rCode, err}
		}
	}
}
<|endoftext|>