<commit_before>package main

import (
    "net/http/httptest"
    "net/http"
    "testing"
    "github.com/stretchr/testify/assert"
    "encoding/json"
    "io/ioutil"
    "fmt"
    "github.com/hashicorp/terraform/helper/schema"
)

type ClientRequest struct {
    GrantType string `json:"grant_type"`
    ClientId string `json:"client_id"`
    ClientSecret string `json:"client_secret"`
    Audience string `json:"audience"`
}

func TestProviderConfigRawSad(t *testing.T) {
    assert := assert.New(t)
    clientSecret := "cauliflower"
    clientId := "joebang"

    testServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.WriteHeader(400)
        w.Header().Set("content-type", "application/json")
        fmt.Fprintf(w, `{"error":"access_denied","error_description":"Service not enabled within domain: https://dshbreak.auth0.com/api/v2/"}`)
    }))
    defer testServer.Close()

    testDomain := testServer.URL[8:]
    result, _ := providerConfigureRaw(testServer.Client(), testDomain, clientId, clientSecret)
    assert.Equal(Config{}, result)

}

func TestProviderConfigRaw(t *testing.T) {
    assert := assert.New(t)

    times := 0
    clientSecret := "cauliflower"
    clientId := "joebang"
    token := "wubbalubbadubdub"

    testServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        times++
        assert.Equal("POST", r.Method)
        body, readErr := ioutil.ReadAll(r.Body)
        if readErr != nil {
            t.Fatalf("Failed to read request: %s", readErr)
        }

        var clientRequest ClientRequest
        unmarshalErr := json.Unmarshal(body, &clientRequest)
        if unmarshalErr != nil {
            t.Fatalf("Failed to parse request: %s", unmarshalErr)
        }

        assert.Equal(clientSecret, clientRequest.ClientSecret)
        assert.Equal(clientId, clientRequest.ClientId)
        assert.Equal("client_credentials", clientRequest.GrantType)

        clientResponse := &Auth0Token{
            AccessToken: token,
            ExpiresIn: 86400,
            Scope: "superman:all",
            TokenType: "type",
        }

        w.WriteHeader(200)
        json.NewEncoder(w).Encode(clientResponse)
    }))
    defer testServer.Close()

    testDomain := testServer.URL[8:]
    result, _ := providerConfigureRaw(testServer.Client(), testDomain, clientId, clientSecret)

    assert.Equal(1, times)
    assert.Equal(testDomain, result.(Config).domain)
    assert.Equal(token, result.(Config).accessToken)
}<commit_msg>Remove unused import.<commit_after>package main

import (
    "net/http/httptest"
    "net/http"
    "testing"
    "github.com/stretchr/testify/assert"
    "encoding/json"
    "io/ioutil"
    "fmt"
)

type ClientRequest struct {
    GrantType string `json:"grant_type"`
    ClientId string `json:"client_id"`
    ClientSecret string `json:"client_secret"`
    Audience string `json:"audience"`
}

func TestProviderConfigRawSad(t *testing.T) {
    assert := assert.New(t)
    clientSecret := "cauliflower"
    clientId := "joebang"

    testServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.WriteHeader(400)
        w.Header().Set("content-type", "application/json")
        fmt.Fprintf(w, `{"error":"access_denied","error_description":"Service not enabled within domain: https://dshbreak.auth0.com/api/v2/"}`)
    }))
    defer testServer.Close()

    testDomain := testServer.URL[8:]
    result, _ := providerConfigureRaw(testServer.Client(), testDomain, clientId, clientSecret)
    assert.Equal(Config{}, result)

}

func TestProviderConfigRaw(t *testing.T) {
    assert := assert.New(t)

    times := 0
    clientSecret := "cauliflower"
    clientId := "joebang"
    token := "wubbalubbadubdub"

    testServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        times++
        assert.Equal("POST", r.Method)
        body, readErr := ioutil.ReadAll(r.Body)
        if readErr != nil {
            t.Fatalf("Failed to read request: %s", readErr)
        }

        var clientRequest ClientRequest
        unmarshalErr := json.Unmarshal(body, &clientRequest)
        if unmarshalErr != nil {
            t.Fatalf("Failed to parse request: %s", unmarshalErr)
        }

        assert.Equal(clientSecret, clientRequest.ClientSecret)
        assert.Equal(clientId, clientRequest.ClientId)
        assert.Equal("client_credentials", clientRequest.GrantType)

        clientResponse := &Auth0Token{
            AccessToken: token,
            ExpiresIn: 86400,
            Scope: "superman:all",
            TokenType: "type",
        }

        w.WriteHeader(200)
        json.NewEncoder(w).Encode(clientResponse)
    }))
    defer testServer.Close()

    testDomain := testServer.URL[8:]
    result, _ := providerConfigureRaw(testServer.Client(), testDomain, clientId, clientSecret)

    assert.Equal(1, times)
    assert.Equal(testDomain, result.(Config).domain)
    assert.Equal(token, result.(Config).accessToken)
}<|endoftext|>


<commit_before>package oauth

import (
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"testing"
)

func TestProviderIsAuthorizedGood(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{})
		c.signer = &MockSigner{}
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	fakeRequest, err := http.NewRequest("GET", "https://example.com/some/path?q=query&q1=another_query", nil)
	if err != nil {
		t.Fatal(err)
	}

	// Set header to good oauth1 header
	fakeRequest.Header.Set(HTTP_AUTH_HEADER, "OAuth oauth_nonce=\"799507437267152061446226936\", oauth_timestamp=\"1446226936\", oauth_version=\"1.0\", oauth_signature_method=\"HMAC-SHA1\", oauth_consumer_key=\"consumerkey\", oauth_signature=\"MOCK_SIGNATURE\"")

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkey", *authorized)
}

func TestProviderIsAuthorizedOauthParamsInBody(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{})
		c.signer = &MockSigner{}
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	fakeRequest, err := http.NewRequest("GET", "https://example.com/some/path?q=query&q1=another_query", nil)
	if err != nil {
		t.Fatal(err)
	}

	// Put good oauth params into the request body
	form := url.Values{}
	form.Set("oauth_consumer_key", "consumerkey")
	form.Set("oauth_nonce", "799507437267152061446226936")
	form.Set("oauth_signature", "MOCK_SIGNATURE")
	form.Set("oauth_signature_method", "HMAC-SHA1")
	form.Set("oauth_timestamp", "1446226936")
	form.Set("oauth_version", "1.0")
	encodedForm := form.Encode()
	fakeRequest.Body = ioutil.NopCloser(strings.NewReader(encodedForm))
	fakeRequest.ContentLength = int64(len(encodedForm))
	fakeRequest.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkey", *authorized)
}

func TestProviderIsAuthorizedOauthParamsInQuery(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{})
		c.signer = &MockSigner{}
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	// Put good oauth params into the query string
	oauthParams := url.Values{}
	oauthParams.Set("oauth_consumer_key", "consumerkey")
	oauthParams.Set("oauth_nonce", "799507437267152061446226936")
	oauthParams.Set("oauth_signature", "MOCK_SIGNATURE")
	oauthParams.Set("oauth_signature_method", "HMAC-SHA1")
	oauthParams.Set("oauth_timestamp", "1446226936")
	oauthParams.Set("oauth_version", "1.0")
	encodedOauthParams := oauthParams.Encode()
	url :=  "https://example.com/some/path?q=query&q1=another_query&" + encodedOauthParams

	fakeRequest, err := http.NewRequest("GET", url, nil)
	if err != nil {
		t.Fatal(err)
	}

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkey", *authorized)
}

func TestProviderIsAuthorizedWithBodyHash(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{BodyHash: true})
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	fakeRequest, err := http.NewRequest("GET", "https://example.com/some/path?q=query&q1=another_query", nil)
	if err != nil {
		t.Fatal(err)
	}

	// Set header to good oauth1 header
	fakeRequest.Header.Set(HTTP_AUTH_HEADER, "OAuth oauth_nonce=\"799507437267152061446226936\", oauth_timestamp=\"1446226936\", oauth_version=\"1.0\", oauth_signature_method=\"HMAC-SHA1\", oauth_consumer_key=\"consumerkey\", oauth_signature=\"RYUiwiUc5LHoipANhDxPbdFHgKc%3D\", oauth_body_hash=\"2jmj7l5rSw0yVb%2FvlWAYkK%2FYBwk%3D\"")

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkey", *authorized)
}

func TestConsumerKeyWithEqualsInIt(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{})
		c.signer = &MockSigner{}
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	fakeRequest, err := http.NewRequest("GET", "https://example.com/some/path?q=query&q1=another_query", nil)
	if err != nil {
		t.Fatal(err)
	}

	// Set header to good oauth1 header
	fakeRequest.Header.Set(HTTP_AUTH_HEADER, "OAuth oauth_nonce=\"799507437267152061446226936\", oauth_timestamp=\"1446226936\", oauth_version=\"1.0\", oauth_signature_method=\"HMAC-SHA1\", oauth_consumer_key=\"consumerkeywithequals=\", oauth_signature=\"MOCK_SIGNATURE\"")

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkeywithequals=", *authorized)
}
<commit_msg>run go fmt, and check in with pre-commit hook<commit_after>package oauth

import (
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"testing"
)

func TestProviderIsAuthorizedGood(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{})
		c.signer = &MockSigner{}
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	fakeRequest, err := http.NewRequest("GET", "https://example.com/some/path?q=query&q1=another_query", nil)
	if err != nil {
		t.Fatal(err)
	}

	// Set header to good oauth1 header
	fakeRequest.Header.Set(HTTP_AUTH_HEADER, "OAuth oauth_nonce=\"799507437267152061446226936\", oauth_timestamp=\"1446226936\", oauth_version=\"1.0\", oauth_signature_method=\"HMAC-SHA1\", oauth_consumer_key=\"consumerkey\", oauth_signature=\"MOCK_SIGNATURE\"")

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkey", *authorized)
}

func TestProviderIsAuthorizedOauthParamsInBody(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{})
		c.signer = &MockSigner{}
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	fakeRequest, err := http.NewRequest("GET", "https://example.com/some/path?q=query&q1=another_query", nil)
	if err != nil {
		t.Fatal(err)
	}

	// Put good oauth params into the request body
	form := url.Values{}
	form.Set("oauth_consumer_key", "consumerkey")
	form.Set("oauth_nonce", "799507437267152061446226936")
	form.Set("oauth_signature", "MOCK_SIGNATURE")
	form.Set("oauth_signature_method", "HMAC-SHA1")
	form.Set("oauth_timestamp", "1446226936")
	form.Set("oauth_version", "1.0")
	encodedForm := form.Encode()
	fakeRequest.Body = ioutil.NopCloser(strings.NewReader(encodedForm))
	fakeRequest.ContentLength = int64(len(encodedForm))
	fakeRequest.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkey", *authorized)
}

func TestProviderIsAuthorizedOauthParamsInQuery(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{})
		c.signer = &MockSigner{}
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	// Put good oauth params into the query string
	oauthParams := url.Values{}
	oauthParams.Set("oauth_consumer_key", "consumerkey")
	oauthParams.Set("oauth_nonce", "799507437267152061446226936")
	oauthParams.Set("oauth_signature", "MOCK_SIGNATURE")
	oauthParams.Set("oauth_signature_method", "HMAC-SHA1")
	oauthParams.Set("oauth_timestamp", "1446226936")
	oauthParams.Set("oauth_version", "1.0")
	encodedOauthParams := oauthParams.Encode()
	url := "https://example.com/some/path?q=query&q1=another_query&" + encodedOauthParams

	fakeRequest, err := http.NewRequest("GET", url, nil)
	if err != nil {
		t.Fatal(err)
	}

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkey", *authorized)
}

func TestProviderIsAuthorizedWithBodyHash(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{BodyHash: true})
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	fakeRequest, err := http.NewRequest("GET", "https://example.com/some/path?q=query&q1=another_query", nil)
	if err != nil {
		t.Fatal(err)
	}

	// Set header to good oauth1 header
	fakeRequest.Header.Set(HTTP_AUTH_HEADER, "OAuth oauth_nonce=\"799507437267152061446226936\", oauth_timestamp=\"1446226936\", oauth_version=\"1.0\", oauth_signature_method=\"HMAC-SHA1\", oauth_consumer_key=\"consumerkey\", oauth_signature=\"RYUiwiUc5LHoipANhDxPbdFHgKc%3D\", oauth_body_hash=\"2jmj7l5rSw0yVb%2FvlWAYkK%2FYBwk%3D\"")

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkey", *authorized)
}

func TestConsumerKeyWithEqualsInIt(t *testing.T) {
	p := NewProvider(func(s string, h map[string]string) (*Consumer, error) {
		c := NewConsumer(s, "consumersecret", ServiceProvider{})
		c.signer = &MockSigner{}
		return c, nil
	})
	p.clock = &MockClock{Time: 1446226936}

	fakeRequest, err := http.NewRequest("GET", "https://example.com/some/path?q=query&q1=another_query", nil)
	if err != nil {
		t.Fatal(err)
	}

	// Set header to good oauth1 header
	fakeRequest.Header.Set(HTTP_AUTH_HEADER, "OAuth oauth_nonce=\"799507437267152061446226936\", oauth_timestamp=\"1446226936\", oauth_version=\"1.0\", oauth_signature_method=\"HMAC-SHA1\", oauth_consumer_key=\"consumerkeywithequals=\", oauth_signature=\"MOCK_SIGNATURE\"")

	authorized, err := p.IsAuthorized(fakeRequest)

	assertEq(t, nil, err)
	assertEq(t, "consumerkeywithequals=", *authorized)
}
<|endoftext|>


<commit_before>// +build !confonly

package dns

import (
	"context"
	"io"
	"sync"

	"golang.org/x/net/dns/dnsmessage"

	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	dns_proto "v2ray.com/core/common/protocol/dns"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/task"
	"v2ray.com/core/features/dns"
	"v2ray.com/core/transport"
	"v2ray.com/core/transport/internet"
)

func init() {
	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		h := new(Handler)
		if err := core.RequireFeatures(ctx, func(dnsClient dns.Client) error {
			return h.Init(config.(*Config), dnsClient)
		}); err != nil {
			return nil, err
		}
		return h, nil
	}))
}

type ownLinkVerifier interface {
	IsOwnLink(ctx context.Context) bool
}

type Handler struct {
	ipv4Lookup      dns.IPv4Lookup
	ipv6Lookup      dns.IPv6Lookup
	ownLinkVerifier ownLinkVerifier
}

func (h *Handler) Init(config *Config, dnsClient dns.Client) error {
	ipv4lookup, ok := dnsClient.(dns.IPv4Lookup)
	if !ok {
		return newError("dns.Client doesn't implement IPv4Lookup")
	}
	h.ipv4Lookup = ipv4lookup

	ipv6lookup, ok := dnsClient.(dns.IPv6Lookup)
	if !ok {
		return newError("dns.Client doesn't implement IPv6Lookup")
	}
	h.ipv6Lookup = ipv6lookup

	if v, ok := dnsClient.(ownLinkVerifier); ok {
		h.ownLinkVerifier = v
	}
	return nil
}

func (h *Handler) isOwnLink(ctx context.Context) bool {
	return h.ownLinkVerifier != nil && h.ownLinkVerifier.IsOwnLink(ctx)
}

func parseIPQuery(b []byte) (r bool, domain string, id uint16, qType dnsmessage.Type) {
	var parser dnsmessage.Parser
	header, err := parser.Start(b)
	if err != nil {
		newError("parser start").Base(err).WriteToLog()
		return
	}

	id = header.ID
	q, err := parser.Question()
	if err != nil {
		newError("question").Base(err).WriteToLog()
		return
	}
	qType = q.Type
	if qType != dnsmessage.TypeA && qType != dnsmessage.TypeAAAA {
		return
	}

	domain = q.Name.String()
	r = true
	return
}

// Process implements proxy.Outbound.
func (h *Handler) Process(ctx context.Context, link *transport.Link, d internet.Dialer) error {
	outbound := session.OutboundFromContext(ctx)
	if outbound == nil || !outbound.Target.IsValid() {
		return newError("invalid outbound")
	}

	dest := outbound.Target

	conn := &outboundConn{
		dialer: func() (internet.Connection, error) {
			return d.Dial(ctx, dest)
		},
		connReady: make(chan struct{}, 1),
	}

	var reader dns_proto.MessageReader
	var writer dns_proto.MessageWriter
	if dest.Network == net.Network_TCP {
		reader = dns_proto.NewTCPReader(link.Reader)
		writer = &dns_proto.TCPWriter{
			Writer: link.Writer,
		}
	} else {
		reader = &dns_proto.UDPReader{
			Reader: link.Reader,
		}
		writer = &dns_proto.UDPWriter{
			Writer: link.Writer,
		}
	}

	var connReader dns_proto.MessageReader
	var connWriter dns_proto.MessageWriter
	if dest.Network == net.Network_TCP {
		connReader = dns_proto.NewTCPReader(buf.NewReader(conn))
		connWriter = &dns_proto.TCPWriter{
			Writer: buf.NewWriter(conn),
		}
	} else {
		connReader = &dns_proto.UDPReader{
			Reader: &buf.PacketReader{Reader: conn},
		}
		connWriter = &dns_proto.UDPWriter{
			Writer: buf.NewWriter(conn),
		}
	}

	request := func() error {
		defer conn.Close()

		for {
			b, err := reader.ReadMessage()
			if err == io.EOF {
				return nil
			}

			if err != nil {
				return err
			}

			if !h.isOwnLink(ctx) {
				isIPQuery, domain, id, qType := parseIPQuery(b.Bytes())
				if isIPQuery {
					go h.handleIPQuery(id, qType, domain, writer)
					continue
				}
			}

			if err := connWriter.WriteMessage(b); err != nil {
				return err
			}
		}
	}

	response := func() error {
		for {
			b, err := connReader.ReadMessage()
			if err == io.EOF {
				return nil
			}

			if err != nil {
				return err
			}

			if err := writer.WriteMessage(b); err != nil {
				return err
			}
		}
	}

	if err := task.Run(ctx, request, response); err != nil {
		return newError("connection ends").Base(err)
	}

	return nil
}

func (h *Handler) handleIPQuery(id uint16, qType dnsmessage.Type, domain string, writer dns_proto.MessageWriter) {
	var ips []net.IP
	var err error

	switch qType {
	case dnsmessage.TypeA:
		ips, err = h.ipv4Lookup.LookupIPv4(domain)
	case dnsmessage.TypeAAAA:
		ips, err = h.ipv6Lookup.LookupIPv6(domain)
	}

	if err != nil {
		newError("ip query").Base(err).WriteToLog()
		return
	}

	if len(ips) == 0 {
		return
	}

	b := buf.New()
	rawBytes := b.Extend(buf.Size)
	builder := dnsmessage.NewBuilder(rawBytes[:0], dnsmessage.Header{
		ID:    id,
		RCode: dnsmessage.RCodeSuccess,
	})
	builder.StartAnswers()

	rHeader := dnsmessage.ResourceHeader{Name: dnsmessage.MustNewName(domain), Class: dnsmessage.ClassINET, TTL: 600}
	for _, ip := range ips {
		if len(ip) == net.IPv4len {
			var r dnsmessage.AResource
			copy(r.A[:], ip)
			builder.AResource(rHeader, r)
		} else {
			var r dnsmessage.AAAAResource
			copy(r.AAAA[:], ip)
			builder.AAAAResource(rHeader, r)
		}
	}
	msgBytes, err := builder.Finish()
	if err != nil {
		newError("pack message").Base(err).WriteToLog()
		b.Release()
		return
	}
	b.Resize(0, int32(len(msgBytes)))

	if err := writer.WriteMessage(b); err != nil {
		newError("write IP answer").Base(err).WriteToLog()
	}
}

type outboundConn struct {
	access sync.Mutex
	dialer func() (internet.Connection, error)

	conn      net.Conn
	connReady chan struct{}
}

func (c *outboundConn) dial() error {
	conn, err := c.dialer()
	if err != nil {
		return err
	}
	c.conn = conn
	c.connReady <- struct{}{}
	return nil
}

func (c *outboundConn) Write(b []byte) (int, error) {
	c.access.Lock()

	if c.conn == nil {
		if err := c.dial(); err != nil {
			c.access.Unlock()
			newError("failed to dial outbound connection").Base(err).AtWarning().WriteToLog()
			return len(b), nil
		}
	}

	c.access.Unlock()

	return c.conn.Write(b)
}

func (c *outboundConn) Read(b []byte) (int, error) {
	var conn net.Conn
	c.access.Lock()
	conn = c.conn
	c.access.Unlock()

	if conn == nil {
		_, open := <-c.connReady
		if !open {
			return 0, io.EOF
		}
		conn = c.conn
	}

	return conn.Read(b)
}

func (c *outboundConn) Close() error {
	c.access.Lock()
	close(c.connReady)
	if c.conn != nil {
		c.conn.Close()
	}
	c.access.Unlock()
	return nil
}
<commit_msg>set response bit in dns<commit_after>// +build !confonly

package dns

import (
	"context"
	"io"
	"sync"

	"golang.org/x/net/dns/dnsmessage"

	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	dns_proto "v2ray.com/core/common/protocol/dns"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/task"
	"v2ray.com/core/features/dns"
	"v2ray.com/core/transport"
	"v2ray.com/core/transport/internet"
)

func init() {
	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		h := new(Handler)
		if err := core.RequireFeatures(ctx, func(dnsClient dns.Client) error {
			return h.Init(config.(*Config), dnsClient)
		}); err != nil {
			return nil, err
		}
		return h, nil
	}))
}

type ownLinkVerifier interface {
	IsOwnLink(ctx context.Context) bool
}

type Handler struct {
	ipv4Lookup      dns.IPv4Lookup
	ipv6Lookup      dns.IPv6Lookup
	ownLinkVerifier ownLinkVerifier
}

func (h *Handler) Init(config *Config, dnsClient dns.Client) error {
	ipv4lookup, ok := dnsClient.(dns.IPv4Lookup)
	if !ok {
		return newError("dns.Client doesn't implement IPv4Lookup")
	}
	h.ipv4Lookup = ipv4lookup

	ipv6lookup, ok := dnsClient.(dns.IPv6Lookup)
	if !ok {
		return newError("dns.Client doesn't implement IPv6Lookup")
	}
	h.ipv6Lookup = ipv6lookup

	if v, ok := dnsClient.(ownLinkVerifier); ok {
		h.ownLinkVerifier = v
	}
	return nil
}

func (h *Handler) isOwnLink(ctx context.Context) bool {
	return h.ownLinkVerifier != nil && h.ownLinkVerifier.IsOwnLink(ctx)
}

func parseIPQuery(b []byte) (r bool, domain string, id uint16, qType dnsmessage.Type) {
	var parser dnsmessage.Parser
	header, err := parser.Start(b)
	if err != nil {
		newError("parser start").Base(err).WriteToLog()
		return
	}

	id = header.ID
	q, err := parser.Question()
	if err != nil {
		newError("question").Base(err).WriteToLog()
		return
	}
	qType = q.Type
	if qType != dnsmessage.TypeA && qType != dnsmessage.TypeAAAA {
		return
	}

	domain = q.Name.String()
	r = true
	return
}

// Process implements proxy.Outbound.
func (h *Handler) Process(ctx context.Context, link *transport.Link, d internet.Dialer) error {
	outbound := session.OutboundFromContext(ctx)
	if outbound == nil || !outbound.Target.IsValid() {
		return newError("invalid outbound")
	}

	dest := outbound.Target

	conn := &outboundConn{
		dialer: func() (internet.Connection, error) {
			return d.Dial(ctx, dest)
		},
		connReady: make(chan struct{}, 1),
	}

	var reader dns_proto.MessageReader
	var writer dns_proto.MessageWriter
	if dest.Network == net.Network_TCP {
		reader = dns_proto.NewTCPReader(link.Reader)
		writer = &dns_proto.TCPWriter{
			Writer: link.Writer,
		}
	} else {
		reader = &dns_proto.UDPReader{
			Reader: link.Reader,
		}
		writer = &dns_proto.UDPWriter{
			Writer: link.Writer,
		}
	}

	var connReader dns_proto.MessageReader
	var connWriter dns_proto.MessageWriter
	if dest.Network == net.Network_TCP {
		connReader = dns_proto.NewTCPReader(buf.NewReader(conn))
		connWriter = &dns_proto.TCPWriter{
			Writer: buf.NewWriter(conn),
		}
	} else {
		connReader = &dns_proto.UDPReader{
			Reader: &buf.PacketReader{Reader: conn},
		}
		connWriter = &dns_proto.UDPWriter{
			Writer: buf.NewWriter(conn),
		}
	}

	request := func() error {
		defer conn.Close()

		for {
			b, err := reader.ReadMessage()
			if err == io.EOF {
				return nil
			}

			if err != nil {
				return err
			}

			if !h.isOwnLink(ctx) {
				isIPQuery, domain, id, qType := parseIPQuery(b.Bytes())
				if isIPQuery {
					go h.handleIPQuery(id, qType, domain, writer)
					continue
				}
			}

			if err := connWriter.WriteMessage(b); err != nil {
				return err
			}
		}
	}

	response := func() error {
		for {
			b, err := connReader.ReadMessage()
			if err == io.EOF {
				return nil
			}

			if err != nil {
				return err
			}

			if err := writer.WriteMessage(b); err != nil {
				return err
			}
		}
	}

	if err := task.Run(ctx, request, response); err != nil {
		return newError("connection ends").Base(err)
	}

	return nil
}

func (h *Handler) handleIPQuery(id uint16, qType dnsmessage.Type, domain string, writer dns_proto.MessageWriter) {
	var ips []net.IP
	var err error

	switch qType {
	case dnsmessage.TypeA:
		ips, err = h.ipv4Lookup.LookupIPv4(domain)
	case dnsmessage.TypeAAAA:
		ips, err = h.ipv6Lookup.LookupIPv6(domain)
	}

	if err != nil {
		newError("ip query").Base(err).WriteToLog()
		return
	}

	if len(ips) == 0 {
		return
	}

	b := buf.New()
	rawBytes := b.Extend(buf.Size)
	builder := dnsmessage.NewBuilder(rawBytes[:0], dnsmessage.Header{
		ID:       id,
		RCode:    dnsmessage.RCodeSuccess,
		Response: true,
	})
	builder.StartAnswers()

	rHeader := dnsmessage.ResourceHeader{Name: dnsmessage.MustNewName(domain), Class: dnsmessage.ClassINET, TTL: 600}
	for _, ip := range ips {
		if len(ip) == net.IPv4len {
			var r dnsmessage.AResource
			copy(r.A[:], ip)
			builder.AResource(rHeader, r)
		} else {
			var r dnsmessage.AAAAResource
			copy(r.AAAA[:], ip)
			builder.AAAAResource(rHeader, r)
		}
	}
	msgBytes, err := builder.Finish()
	if err != nil {
		newError("pack message").Base(err).WriteToLog()
		b.Release()
		return
	}
	b.Resize(0, int32(len(msgBytes)))

	if err := writer.WriteMessage(b); err != nil {
		newError("write IP answer").Base(err).WriteToLog()
	}
}

type outboundConn struct {
	access sync.Mutex
	dialer func() (internet.Connection, error)

	conn      net.Conn
	connReady chan struct{}
}

func (c *outboundConn) dial() error {
	conn, err := c.dialer()
	if err != nil {
		return err
	}
	c.conn = conn
	c.connReady <- struct{}{}
	return nil
}

func (c *outboundConn) Write(b []byte) (int, error) {
	c.access.Lock()

	if c.conn == nil {
		if err := c.dial(); err != nil {
			c.access.Unlock()
			newError("failed to dial outbound connection").Base(err).AtWarning().WriteToLog()
			return len(b), nil
		}
	}

	c.access.Unlock()

	return c.conn.Write(b)
}

func (c *outboundConn) Read(b []byte) (int, error) {
	var conn net.Conn
	c.access.Lock()
	conn = c.conn
	c.access.Unlock()

	if conn == nil {
		_, open := <-c.connReady
		if !open {
			return 0, io.EOF
		}
		conn = c.conn
	}

	return conn.Read(b)
}

func (c *outboundConn) Close() error {
	c.access.Lock()
	close(c.connReady)
	if c.conn != nil {
		c.conn.Close()
	}
	c.access.Unlock()
	return nil
}
<|endoftext|>


<commit_before>package frame

import (
	"errors"
	"fmt"
	dbg "fmt"
	"math"

	"github.com/eaburns/bit"
)

// A SubFrame contains the decoded audio data of a channel.
type SubFrame struct {
	// Header specifies the attributes of the subframe, like prediction method
	// and order, residual coding parameters, etc.
	Header *SubHeader
	// Samples contains the decoded audio samples of the channel.
	Samples []Sample
}

// A Sample is an audio sample. The size of each sample is between 4 and 32 bits.
type Sample uint32

// NewSubFrame parses and returns a new subframe, which consists of a subframe
// header and encoded audio samples.
//
// Subframe format (pseudo code):
//
//    type SUBFRAME struct {
//       header      SUBFRAME_HEADER
//       enc_samples SUBFRAME_CONSTANT || SUBFRAME_FIXED || SUBFRAME_LPC || SUBFRAME_VERBATIM
//    }
//
// ref: http://flac.sourceforge.net/format.html#subframe
func (h *Header) NewSubFrame(br *bit.Reader) (subframe *SubFrame, err error) {
	// Parse subframe header.
	subframe = new(SubFrame)
	subframe.Header, err = h.NewSubHeader(br)
	if err != nil {
		return nil, err
	}

	// Decode samples.
	sh := subframe.Header
	switch sh.PredMethod {
	case PredConstant:
		subframe.Samples, err = h.DecodeConstant(br)
	case PredFixed:
		subframe.Samples, err = h.DecodeFixed(br, int(sh.PredOrder))
	case PredLPC:
		subframe.Samples, err = h.DecodeLPC(br, int(sh.PredOrder))
	case PredVerbatim:
		subframe.Samples, err = h.DecodeVerbatim(br)
	default:
		return nil, fmt.Errorf("Header.NewSubFrame: unknown subframe prediction method: %d", sh.PredMethod)
	}
	if err != nil {
		return nil, err
	}

	return subframe, nil
}

// A SubHeader is a subframe header, which contains information about how the
// subframe audio samples are encoded.
type SubHeader struct {
	// PredMethod is the subframe prediction method.
	PredMethod PredMethod
	// WastedBitCount is the number of wasted bits per sample.
	WastedBitCount int8
	// PredOrder is the subframe predictor order, which is used accordingly:
	//    Fixed: Predictor order.
	//    LPC:   LPC order.
	PredOrder int8
}

// PredMethod specifies the subframe prediction method.
type PredMethod int8

// Subframe prediction methods.
const (
	PredConstant PredMethod = iota
	PredFixed
	PredLPC
	PredVerbatim
)

// NewSubHeader parses and returns a new subframe header.
//
// Subframe header format (pseudo code):
//    type SUBFRAME_HEADER struct {
//       _                uint1 // zero-padding, to prevent sync-fooling.
//       type             uint6
//       // 0: no wasted bits-per-sample in source subblock, k = 0.
//       // 1: k wasted bits-per-sample in source subblock, k-1 follows, unary
//       // coded; e.g. k=3 => 001 follows, k=7 => 0000001 follows.
//       wasted_bit_count uint1+k
//    }
//
// ref: http://flac.sourceforge.net/format.html#subframe_header
func (h *Header) NewSubHeader(br *bit.Reader) (sh *SubHeader, err error) {
	// field 0: padding (1 bit)
	// field 1: type    (6 bits)
	fields, err := br.ReadFields(1, 6)
	if err != nil {
		return nil, err
	}

	// Padding.
	if fields[0] != 0 {
		return nil, errors.New("Header.NewSubHeader: invalid padding; must be 0")
	}

	// Subframe prediction method.
	//    000000: SUBFRAME_CONSTANT
	//    000001: SUBFRAME_VERBATIM
	//    00001x: reserved
	//    0001xx: reserved
	//    001xxx: if(xxx <= 4) SUBFRAME_FIXED, xxx=order ; else reserved
	//    01xxxx: reserved
	//    1xxxxx: SUBFRAME_LPC, xxxxx=order-1
	sh = new(SubHeader)
	n := fields[1]
	switch {
	case n == 0:
		// 000000: SUBFRAME_CONSTANT
		sh.PredMethod = PredConstant
	case n == 1:
		// 000001: SUBFRAME_VERBATIM
		sh.PredMethod = PredVerbatim
	case n < 8:
		// 00001x: reserved
		// 0001xx: reserved
		return nil, fmt.Errorf("Header.NewSubHeader: invalid subframe prediction method; reserved bit pattern: %06b", n)
	case n < 16:
		// 001xxx: if(xxx <= 4) SUBFRAME_FIXED, xxx=order ; else reserved
		const predOrderMask = 0x07
		sh.PredOrder = int8(n) & predOrderMask
		if sh.PredOrder > 4 {
			return nil, fmt.Errorf("Header.NewSubHeader: invalid subframe prediction method; reserved bit pattern: %06b", n)
		}
		sh.PredMethod = PredFixed
	case n < 32:
		// 01xxxx: reserved
		return nil, fmt.Errorf("Header.NewSubHeader: invalid subframe prediction method; reserved bit pattern: %06b", n)
	case n < 64:
		// 1xxxxx: SUBFRAME_LPC, xxxxx=order-1
		const predOrderMask = 0x1F
		sh.PredOrder = int8(n)&predOrderMask + 1
		sh.PredMethod = PredLPC
	default:
		// should be unreachable.
		return nil, fmt.Errorf("Header.NewSubHeader: unhandled subframe prediction method; bit pattern: %06b", n)
	}

	// Wasted bits-per-sample, 1+k bits.
	bits, err := br.Read(1)
	if err != nil {
		return nil, err
	}
	if bits != 0 {
		// k wasted bits-per-sample in source subblock, k-1 follows, unary coded;
		// e.g. k=3 => 001 follows, k=7 => 0000001 follows.
		/// ### [ todo ] ###
		///    - verify.
		/// ### [/ todo ] ###
		for {
			sh.WastedBitCount++
			bits, err = br.Read(1)
			if err != nil {
				return nil, err
			}
			if bits == 1 {
				break
			}
		}
	}

	return sh, nil
}

// DecodeConstant decodes and returns a slice of samples. The first sample is
// constant throughout the entire subframe.
//
// ref: http://flac.sourceforge.net/format.html#subframe_constant
func (h *Header) DecodeConstant(br *bit.Reader) (samples []Sample, err error) {
	// Read constant sample.
	bits, err := br.Read(uint(h.SampleSize))
	if err != nil {
		return nil, err
	}
	sample := Sample(bits)
	dbg.Println("Constant sample:", sample)

	// Duplicate the constant sample, sample count number of times.
	for i := uint16(0); i < h.SampleCount; i++ {
		samples = append(samples, sample)
	}

	return samples, nil
}

// DecodeFixed decodes and returns a slice of samples.
/// ### [ todo ] ###
///    - add more details.
/// ### [/ todo ] ###
//
// ref: http://flac.sourceforge.net/format.html#subframe_fixed
func (h *Header) DecodeFixed(br *bit.Reader, predOrder int) (samples []Sample, err error) {
	// Unencoded warm-up samples:
	//    n bits = frame's bits-per-sample * predictor order
	for i := 0; i < predOrder; i++ {
		bits, err := br.Read(uint(h.SampleSize))
		if err != nil {
			return nil, err
		}
		sample := Sample(bits)
		dbg.Println("Fixed warm-up sample:", sample)
		samples = append(samples, sample)
	}

	residuals, err := h.DecodeResidual(br, predOrder)
	if err != nil {
		return nil, err
	}
	_ = residuals
	/// ### [ todo ] ###
	///    - not yet implemented.
	/// ### [/ todo ] ###
	return nil, errors.New("not yet implemented; Fixed encoding")
}

// DecodeLPC decodes and returns a slice of samples.
/// ### [ todo ] ###
///    - add more details.
/// ### [/ todo ] ###
//
// ref: http://flac.sourceforge.net/format.html#subframe_lpc
func (h *Header) DecodeLPC(br *bit.Reader, lpcOrder int) (samples []Sample, err error) {
	// Unencoded warm-up samples:
	//    n bits = frame's bits-per-sample * lpc order
	for i := 0; i < lpcOrder; i++ {
		bits, err := br.Read(uint(h.SampleSize))
		if err != nil {
			return nil, err
		}
		sample := Sample(bits)
		dbg.Println("LPC warm-up sample:", sample)
		samples = append(samples, sample)
	}

	// (Quantized linear predictor coefficients' precision in bits) - 1.
	n, err := br.Read(4)
	if err != nil {
		return nil, err
	}
	if n == 0x0F {
		// 1111: invalid.
		return nil, errors.New("Header.DecodeLPC: invalid quantized lpc precision; reserved bit pattern: 1111")
	}
	qlpcPrec := int(n) + 1

	// Quantized linear predictor coefficient shift needed in bits.
	qlpcShift, err := br.Read(5)
	if err != nil {
		return nil, err
	}
	/// ### [ todo ] ###
	///    - NOTE: this number is signed two's-complement.
	///    - special case for negative numbers required?
	///    - the same goes for qlpcPrec.
	/// ### [/ todo ] ###
	_ = qlpcShift

	// Unencoded predictor coefficients.
	for i := 0; i < lpcOrder; i++ {
		pc, err := br.Read(uint(qlpcPrec))
		if err != nil {
			return nil, err
		}
		dbg.Println("pc:", pc)
		_ = pc
	}

	residuals, err := h.DecodeResidual(br, lpcOrder)
	if err != nil {
		return nil, err
	}
	_ = residuals
	/// ### [ todo ] ###
	///    - not yet implemented.
	/// ### [/ todo ] ###
	return nil, errors.New("not yet implemented; LPC encoding")
}

// DecodeVerbatim decodes and returns a slice of samples. The samples are stored
// unencoded.
//
// ref: http://flac.sourceforge.net/format.html#subframe_verbatim
func (h *Header) DecodeVerbatim(br *bit.Reader) (samples []Sample, err error) {
	// Read unencoded samples.
	for i := uint16(0); i < h.SampleCount; i++ {
		bits, err := br.Read(uint(h.SampleSize))
		if err != nil {
			return nil, err
		}
		sample := Sample(bits)
		dbg.Println("Verbatim sample:", sample)
		samples = append(samples, sample)
	}

	return samples, nil
}

// DecodeResidual decodes and returns a slice of residuals.
/// ### [ todo ] ###
///    - add more details.
/// ### [/ todo ] ###
//
// ref: http://flac.sourceforge.net/format.html#residual
func (h *Header) DecodeResidual(br *bit.Reader, predOrder int) (residuals []int, err error) {
	// Residual coding method.
	method, err := br.Read(2)
	if err != nil {
		return nil, err
	}
	switch method {
	case 0:
		// 00: partitioned Rice coding with 4-bit Rice parameter;
		//     RESIDUAL_CODING_METHOD_PARTITIONED_RICE follows
		return h.DecodeRice(br, predOrder)
	case 1:
		// 01: partitioned Rice coding with 5-bit Rice parameter;
		//     RESIDUAL_CODING_METHOD_PARTITIONED_RICE2 follows
		return h.DecodeRice2(br, predOrder)
	}
	// 1x: reserved
	return nil, fmt.Errorf("Header.DecodeResidual: invalid residual coding method; reserved bit pattern: %02b", method)
}

// DecodeRice decodes and returns a slice of residuals. The residual coding
// method used is partitioned Rice coding with a 4-bit Rice parameter.
//
// ref: http://flac.sourceforge.net/format.html#partitioned_rice
func (h *Header) DecodeRice(br *bit.Reader, predOrder int) (residuals []int, err error) {
	// Partition order.
	partOrder, err := br.Read(4)
	if err != nil {
		return nil, err
	}

	// Rice partitions.
	partCount := int(math.Pow(2, float64(partOrder)))
	partSampleCount := int(h.SampleCount) / partCount
	for partNum := 0; partNum < partCount; partNum++ {
		// Encoding parameter.
		n, err := br.Read(4)
		if err != nil {
			return nil, err
		}
		if n == 0x0F {
			// 1111: Escape code, meaning the partition is in unencoded binary form
			// using n bits per sample; n follows as a 5-bit number.
			/// ### [ todo ] ###
			///    - not yet implemented.
			/// ### [/ todo ] ###
			return nil, errors.New("not yet implemented; rice encoding parameter escape code")
		}
		riceParam := n
		_ = riceParam

		dbg.Println("riceParam:", riceParam)

		// Encoded residual.
		sampleCount := partSampleCount
		if partNum == 0 {
			sampleCount -= predOrder
		}

		dbg.Println("sampleCount:", sampleCount)
	}

	/// ### [ todo ] ###
	///    - not yet implemented.
	/// ### [/ todo ] ###
	return nil, errors.New("not yet implemented; rice coding method 0")
}

// DecodeRice2 decodes and returns a slice of residuals. The residual coding
// method used is partitioned Rice coding with a 5-bit Rice parameter.
//
// ref: http://flac.sourceforge.net/format.html#partitioned_rice2
func (h *Header) DecodeRice2(br *bit.Reader, predOrder int) (residuals []int, err error) {
	/// ### [ todo ] ###
	///    - not yet implemented.
	/// ### [/ todo ] ###
	return nil, errors.New("not yet implemented; rice coding method 1")
}

/**
type SubFrameLpc struct {
	Precision             uint8
	ShiftNeeded           uint8
	PredictorCoefficients []byte
}

type Rice struct {
	Partitions     []RicePartition
}

type Rice2 struct {
	Partitions     []Rice2Partition
}

type RicePartition struct {
	EncodingParameter uint16
}

type Rice2Partition struct{}
*/
<commit_msg>frame: Add some comments and panic on unreachable branches.<commit_after>package frame

import (
	"errors"
	"fmt"
	dbg "fmt"
	"math"

	"github.com/eaburns/bit"
)

// A SubFrame contains the decoded audio data of a channel.
type SubFrame struct {
	// Header specifies the attributes of the subframe, like prediction method
	// and order, residual coding parameters, etc.
	Header *SubHeader
	// Samples contains the decoded audio samples of the channel.
	Samples []Sample
}

// A Sample is an audio sample. The size of each sample is between 4 and 32 bits.
type Sample uint32

// NewSubFrame parses and returns a new subframe, which consists of a subframe
// header and encoded audio samples.
//
// Subframe format (pseudo code):
//
//    type SUBFRAME struct {
//       header      SUBFRAME_HEADER
//       enc_samples SUBFRAME_CONSTANT || SUBFRAME_FIXED || SUBFRAME_LPC || SUBFRAME_VERBATIM
//    }
//
// ref: http://flac.sourceforge.net/format.html#subframe
func (h *Header) NewSubFrame(br *bit.Reader) (subframe *SubFrame, err error) {
	// Parse subframe header.
	subframe = new(SubFrame)
	subframe.Header, err = h.NewSubHeader(br)
	if err != nil {
		return nil, err
	}

	// Decode samples.
	sh := subframe.Header
	switch sh.PredMethod {
	case PredConstant:
		subframe.Samples, err = h.DecodeConstant(br)
	case PredFixed:
		subframe.Samples, err = h.DecodeFixed(br, int(sh.PredOrder))
	case PredLPC:
		subframe.Samples, err = h.DecodeLPC(br, int(sh.PredOrder))
	case PredVerbatim:
		subframe.Samples, err = h.DecodeVerbatim(br)
	default:
		return nil, fmt.Errorf("Header.NewSubFrame: unknown subframe prediction method: %d", sh.PredMethod)
	}
	if err != nil {
		return nil, err
	}

	return subframe, nil
}

// A SubHeader is a subframe header, which contains information about how the
// subframe audio samples are encoded.
type SubHeader struct {
	// PredMethod is the subframe prediction method.
	PredMethod PredMethod
	// WastedBitCount is the number of wasted bits per sample.
	WastedBitCount int8
	// PredOrder is the subframe predictor order, which is used accordingly:
	//    Fixed: Predictor order.
	//    LPC:   LPC order.
	PredOrder int8
}

// PredMethod specifies the subframe prediction method.
type PredMethod int8

// Subframe prediction methods.
const (
	PredConstant PredMethod = iota
	PredFixed
	PredLPC
	PredVerbatim
)

// NewSubHeader parses and returns a new subframe header.
//
// Subframe header format (pseudo code):
//    type SUBFRAME_HEADER struct {
//       _                uint1 // zero-padding, to prevent sync-fooling.
//       type             uint6
//       // 0: no wasted bits-per-sample in source subblock, k = 0.
//       // 1: k wasted bits-per-sample in source subblock, k-1 follows, unary
//       // coded; e.g. k=3 => 001 follows, k=7 => 0000001 follows.
//       wasted_bit_count uint1+k
//    }
//
// ref: http://flac.sourceforge.net/format.html#subframe_header
func (h *Header) NewSubHeader(br *bit.Reader) (sh *SubHeader, err error) {
	// field 0: padding (1 bit)
	// field 1: type    (6 bits)
	fields, err := br.ReadFields(1, 6)
	if err != nil {
		return nil, err
	}

	// Padding.
	// field 0: padding (1 bit)
	if fields[0] != 0 {
		return nil, errors.New("Header.NewSubHeader: invalid padding; must be 0")
	}

	// Subframe prediction method.
	//    000000: SUBFRAME_CONSTANT
	//    000001: SUBFRAME_VERBATIM
	//    00001x: reserved
	//    0001xx: reserved
	//    001xxx: if(xxx <= 4) SUBFRAME_FIXED, xxx=order ; else reserved
	//    01xxxx: reserved
	//    1xxxxx: SUBFRAME_LPC, xxxxx=order-1
	sh = new(SubHeader)
	// field 1: type (6 bits)
	n := fields[1]
	switch {
	case n == 0:
		// 000000: SUBFRAME_CONSTANT
		sh.PredMethod = PredConstant
	case n == 1:
		// 000001: SUBFRAME_VERBATIM
		sh.PredMethod = PredVerbatim
	case n < 8:
		// 00001x: reserved
		// 0001xx: reserved
		return nil, fmt.Errorf("Header.NewSubHeader: invalid subframe prediction method; reserved bit pattern: %06b", n)
	case n < 16:
		// 001xxx: if(xxx <= 4) SUBFRAME_FIXED, xxx=order ; else reserved
		const predOrderMask = 0x07
		sh.PredOrder = int8(n) & predOrderMask
		if sh.PredOrder > 4 {
			return nil, fmt.Errorf("Header.NewSubHeader: invalid subframe prediction method; reserved bit pattern: %06b", n)
		}
		sh.PredMethod = PredFixed
	case n < 32:
		// 01xxxx: reserved
		return nil, fmt.Errorf("Header.NewSubHeader: invalid subframe prediction method; reserved bit pattern: %06b", n)
	case n < 64:
		// 1xxxxx: SUBFRAME_LPC, xxxxx=order-1
		const predOrderMask = 0x1F
		sh.PredOrder = int8(n)&predOrderMask + 1
		sh.PredMethod = PredLPC
	default:
		// should be unreachable.
		panic(fmt.Errorf("Header.NewSubHeader: unhandled subframe prediction method; bit pattern: %06b", n))
	}

	// Wasted bits-per-sample, 1+k bits.
	bits, err := br.Read(1)
	if err != nil {
		return nil, err
	}
	if bits != 0 {
		// k wasted bits-per-sample in source subblock, k-1 follows, unary coded;
		// e.g. k=3 => 001 follows, k=7 => 0000001 follows.
		/// ### [ todo ] ###
		///    - verify.
		/// ### [/ todo ] ###
		for {
			sh.WastedBitCount++
			bits, err = br.Read(1)
			if err != nil {
				return nil, err
			}
			if bits == 1 {
				break
			}
		}
	}

	return sh, nil
}

// DecodeConstant decodes and returns a slice of samples. The first sample is
// constant throughout the entire subframe.
//
// ref: http://flac.sourceforge.net/format.html#subframe_constant
func (h *Header) DecodeConstant(br *bit.Reader) (samples []Sample, err error) {
	// Read constant sample.
	bits, err := br.Read(uint(h.SampleSize))
	if err != nil {
		return nil, err
	}
	sample := Sample(bits)
	dbg.Println("Constant sample:", sample)

	// Duplicate the constant sample, sample count number of times.
	for i := uint16(0); i < h.SampleCount; i++ {
		samples = append(samples, sample)
	}

	return samples, nil
}

// DecodeFixed decodes and returns a slice of samples.
/// ### [ todo ] ###
///    - add more details.
/// ### [/ todo ] ###
//
// ref: http://flac.sourceforge.net/format.html#subframe_fixed
func (h *Header) DecodeFixed(br *bit.Reader, predOrder int) (samples []Sample, err error) {
	// Unencoded warm-up samples:
	//    n bits = frame's bits-per-sample * predictor order
	for i := 0; i < predOrder; i++ {
		bits, err := br.Read(uint(h.SampleSize))
		if err != nil {
			return nil, err
		}
		sample := Sample(bits)
		dbg.Println("Fixed warm-up sample:", sample)
		samples = append(samples, sample)
	}

	residuals, err := h.DecodeResidual(br, predOrder)
	if err != nil {
		return nil, err
	}
	_ = residuals
	/// ### [ todo ] ###
	///    - not yet implemented.
	/// ### [/ todo ] ###
	return nil, errors.New("not yet implemented; Fixed encoding")
}

// DecodeLPC decodes and returns a slice of samples.
/// ### [ todo ] ###
///    - add more details.
/// ### [/ todo ] ###
//
// ref: http://flac.sourceforge.net/format.html#subframe_lpc
func (h *Header) DecodeLPC(br *bit.Reader, lpcOrder int) (samples []Sample, err error) {
	// Unencoded warm-up samples:
	//    n bits = frame's bits-per-sample * lpc order
	for i := 0; i < lpcOrder; i++ {
		bits, err := br.Read(uint(h.SampleSize))
		if err != nil {
			return nil, err
		}
		sample := Sample(bits)
		dbg.Println("LPC warm-up sample:", sample)
		samples = append(samples, sample)
	}

	// (Quantized linear predictor coefficients' precision in bits) - 1.
	n, err := br.Read(4)
	if err != nil {
		return nil, err
	}
	if n == 0x0F {
		// 1111: invalid.
		return nil, errors.New("Header.DecodeLPC: invalid quantized lpc precision; reserved bit pattern: 1111")
	}
	qlpcPrec := int(n) + 1

	// Quantized linear predictor coefficient shift needed in bits.
	qlpcShift, err := br.Read(5)
	if err != nil {
		return nil, err
	}
	/// ### [ todo ] ###
	///    - NOTE: this number is signed two's-complement.
	///    - special case for negative numbers required?
	///    - the same goes for qlpcPrec.
	/// ### [/ todo ] ###
	_ = qlpcShift

	// Unencoded predictor coefficients.
	for i := 0; i < lpcOrder; i++ {
		pc, err := br.Read(uint(qlpcPrec))
		if err != nil {
			return nil, err
		}
		dbg.Println("pc:", pc)
		_ = pc
	}

	residuals, err := h.DecodeResidual(br, lpcOrder)
	if err != nil {
		return nil, err
	}
	_ = residuals
	/// ### [ todo ] ###
	///    - not yet implemented.
	/// ### [/ todo ] ###
	return nil, errors.New("not yet implemented; LPC encoding")
}

// DecodeVerbatim decodes and returns a slice of samples. The samples are stored
// unencoded.
//
// ref: http://flac.sourceforge.net/format.html#subframe_verbatim
func (h *Header) DecodeVerbatim(br *bit.Reader) (samples []Sample, err error) {
	// Read unencoded samples.
	for i := uint16(0); i < h.SampleCount; i++ {
		bits, err := br.Read(uint(h.SampleSize))
		if err != nil {
			return nil, err
		}
		sample := Sample(bits)
		dbg.Println("Verbatim sample:", sample)
		samples = append(samples, sample)
	}

	return samples, nil
}

// DecodeResidual decodes and returns a slice of residuals.
/// ### [ todo ] ###
///    - add more details.
/// ### [/ todo ] ###
//
// ref: http://flac.sourceforge.net/format.html#residual
func (h *Header) DecodeResidual(br *bit.Reader, predOrder int) (residuals []int, err error) {
	// Residual coding method.
	method, err := br.Read(2)
	if err != nil {
		return nil, err
	}
	switch method {
	case 0:
		// 00: partitioned Rice coding with 4-bit Rice parameter;
		//     RESIDUAL_CODING_METHOD_PARTITIONED_RICE follows
		return h.DecodeRice(br, predOrder)
	case 1:
		// 01: partitioned Rice coding with 5-bit Rice parameter;
		//     RESIDUAL_CODING_METHOD_PARTITIONED_RICE2 follows
		return h.DecodeRice2(br, predOrder)
	}
	// 1x: reserved
	return nil, fmt.Errorf("Header.DecodeResidual: invalid residual coding method; reserved bit pattern: %02b", method)
}

// DecodeRice decodes and returns a slice of residuals. The residual coding
// method used is partitioned Rice coding with a 4-bit Rice parameter.
//
// ref: http://flac.sourceforge.net/format.html#partitioned_rice
func (h *Header) DecodeRice(br *bit.Reader, predOrder int) (residuals []int, err error) {
	// Partition order.
	partOrder, err := br.Read(4)
	if err != nil {
		return nil, err
	}

	// Rice partitions.
	partCount := int(math.Pow(2, float64(partOrder)))
	partSampleCount := int(h.SampleCount) / partCount
	for partNum := 0; partNum < partCount; partNum++ {
		// Encoding parameter.
		n, err := br.Read(4)
		if err != nil {
			return nil, err
		}
		if n == 0x0F {
			// 1111: Escape code, meaning the partition is in unencoded binary form
			// using n bits per sample; n follows as a 5-bit number.
			/// ### [ todo ] ###
			///    - not yet implemented.
			/// ### [/ todo ] ###
			return nil, errors.New("not yet implemented; rice encoding parameter escape code")
		}
		riceParam := n
		_ = riceParam

		dbg.Println("riceParam:", riceParam)

		// Encoded residual.
		sampleCount := partSampleCount
		if partNum == 0 {
			sampleCount -= predOrder
		}

		dbg.Println("sampleCount:", sampleCount)
	}

	/// ### [ todo ] ###
	///    - not yet implemented.
	/// ### [/ todo ] ###
	return nil, errors.New("not yet implemented; rice coding method 0")
}

// DecodeRice2 decodes and returns a slice of residuals. The residual coding
// method used is partitioned Rice coding with a 5-bit Rice parameter.
//
// ref: http://flac.sourceforge.net/format.html#partitioned_rice2
func (h *Header) DecodeRice2(br *bit.Reader, predOrder int) (residuals []int, err error) {
	/// ### [ todo ] ###
	///    - not yet implemented.
	/// ### [/ todo ] ###
	return nil, errors.New("not yet implemented; rice coding method 1")
}

/**
type SubFrameLpc struct {
	Precision             uint8
	ShiftNeeded           uint8
	PredictorCoefficients []byte
}

type Rice struct {
	Partitions     []RicePartition
}

type Rice2 struct {
	Partitions     []Rice2Partition
}

type RicePartition struct {
	EncodingParameter uint16
}

type Rice2Partition struct{}
*/
<|endoftext|>


<commit_before>package domain

import (
	"fmt"
	"net/http"
	"reflect"
	"time"

	"sync"

	"github.com/pivotal-golang/lager"
)

//go:generate counterfeiter . UrlGetter
type UrlGetter interface {
	Get(url string) (*http.Response, error)
}

func HttpUrlGetterProvider(healthcheckTimeout time.Duration) UrlGetter {
	return &http.Client{
		Timeout: healthcheckTimeout,
	}
}

var UrlGetterProvider = HttpUrlGetterProvider

type Cluster struct {
	mutex              sync.RWMutex
	backends           Backends
	logger             lager.Logger
	healthcheckTimeout time.Duration
	arpManager         ArpManager
	message            string
	lastUpdated        time.Time
}

func NewCluster(backends Backends, healthcheckTimeout time.Duration, logger lager.Logger, arpManager ArpManager) *Cluster {
	return &Cluster{
		backends:           backends,
		logger:             logger,
		healthcheckTimeout: healthcheckTimeout,
		arpManager:         arpManager,
	}
}

func (c *Cluster) Monitor(stopChan <-chan interface{}) {
	client := UrlGetterProvider(c.healthcheckTimeout)

	for b := range c.backends.All() {
		go func(backend Backend) {
			counters := c.setupCounters()
			for {
				select {
				case <-time.After(c.healthcheckTimeout / 5):

					counters.IncrementCount("dial")
					shouldLog := counters.Should("log")

					url := backend.HealthcheckUrl()
					resp, err := client.Get(url)

					if err == nil && resp.StatusCode == http.StatusOK {
						c.backends.SetHealthy(backend)
						counters.ResetCount("consecutiveUnhealthyChecks")

						if shouldLog {
							c.logger.Debug("Healthcheck succeeded", lager.Data{"endpoint": url})
						}
					} else {
						c.backends.SetUnhealthy(backend)
						counters.IncrementCount("consecutiveUnhealthyChecks")

						if shouldLog {
							c.logger.Error(
								"Healthcheck failed on backend",
								fmt.Errorf("Non-200 status code from healthcheck"),
								lager.Data{
									"backend":  backend.AsJSON(),
									"endpoint": url,
									"resp":     fmt.Sprintf("%#v", resp),
									"err":      err,
								},
							)
						}
					}

					if counters.Should("clearArp") {
						backendHost := backend.AsJSON().Host

						if c.arpManager.IsCached(backendHost) {
							err = c.arpManager.ClearCache(backendHost)
							if err != nil {
								c.logger.Error("Failed to clear arp cache", err)
							}
						}
					}
				case <-stopChan:
					return
				}
			}
		}(b)
	}
}

func (c *Cluster) setupCounters() *DecisionCounters {
	counters := NewDecisionCounters()
	logFreq := uint64(5)
	clearArpFreq := uint64(5)

	//used to make logs less noisy
	counters.AddCondition("log", func() bool {
		return (counters.GetCount("dial") % logFreq) == 0
	})

	//only clear ARP cache after X consecutive unhealthy dials
	counters.AddCondition("clearArp", func() bool {
		// golang makes it difficult to tell whether the value of an interface is nil
		if reflect.ValueOf(c.arpManager).IsNil() {
			return false
		} else {
			checks := counters.GetCount("consecutiveUnhealthyChecks")
			return (checks > 0) && (checks%clearArpFreq) == 0
		}
	})

	return counters
}
<commit_msg>WIP: ClusterMonitor is smaller<commit_after>package domain

import (
	"fmt"
	"net/http"
	"reflect"
	"time"

	"github.com/pivotal-golang/lager"
)

//go:generate counterfeiter . UrlGetter
type UrlGetter interface {
	Get(url string) (*http.Response, error)
}

func HttpUrlGetterProvider(healthcheckTimeout time.Duration) UrlGetter {
	return &http.Client{
		Timeout: healthcheckTimeout,
	}
}

var UrlGetterProvider = HttpUrlGetterProvider

type Cluster struct {
	backends           Backends
	logger             lager.Logger
	healthcheckTimeout time.Duration
	arpManager         ArpManager
}

func NewCluster(backends Backends, healthcheckTimeout time.Duration, logger lager.Logger, arpManager ArpManager) *Cluster {
	return &Cluster{
		backends:           backends,
		logger:             logger,
		healthcheckTimeout: healthcheckTimeout,
		arpManager:         arpManager,
	}
}

func (c *Cluster) Monitor(stopChan <-chan interface{}) {
	client := UrlGetterProvider(c.healthcheckTimeout)

	for b := range c.backends.All() {
		go func(backend Backend) {
			counters := c.setupCounters()
			for {
				select {
				case <-time.After(c.healthcheckTimeout / 5):

					counters.IncrementCount("dial")
					shouldLog := counters.Should("log")

					url := backend.HealthcheckUrl()
					resp, err := client.Get(url)

					if err == nil && resp.StatusCode == http.StatusOK {
						c.backends.SetHealthy(backend)
						counters.ResetCount("consecutiveUnhealthyChecks")

						if shouldLog {
							c.logger.Debug("Healthcheck succeeded", lager.Data{"endpoint": url})
						}
					} else {
						c.backends.SetUnhealthy(backend)
						counters.IncrementCount("consecutiveUnhealthyChecks")

						if shouldLog {
							c.logger.Error(
								"Healthcheck failed on backend",
								fmt.Errorf("Non-200 status code from healthcheck"),
								lager.Data{
									"backend":  backend.AsJSON(),
									"endpoint": url,
									"resp":     fmt.Sprintf("%#v", resp),
									"err":      err,
								},
							)
						}
					}

					if counters.Should("clearArp") {
						backendHost := backend.AsJSON().Host

						if c.arpManager.IsCached(backendHost) {
							err = c.arpManager.ClearCache(backendHost)
							if err != nil {
								c.logger.Error("Failed to clear arp cache", err)
							}
						}
					}
				case <-stopChan:
					return
				}
			}
		}(b)
	}
}

func (c *Cluster) setupCounters() *DecisionCounters {
	counters := NewDecisionCounters()
	logFreq := uint64(5)
	clearArpFreq := uint64(5)

	//used to make logs less noisy
	counters.AddCondition("log", func() bool {
		return (counters.GetCount("dial") % logFreq) == 0
	})

	//only clear ARP cache after X consecutive unhealthy dials
	counters.AddCondition("clearArp", func() bool {
		// golang makes it difficult to tell whether the value of an interface is nil
		if reflect.ValueOf(c.arpManager).IsNil() {
			return false
		} else {
			checks := counters.GetCount("consecutiveUnhealthyChecks")
			return (checks > 0) && (checks%clearArpFreq) == 0
		}
	})

	return counters
}
<|endoftext|>


<commit_before>package goflow

import (
	"testing"
)

type withInvalidPorts struct {
	NotChan int
	Chan    <-chan int
}

func (c *withInvalidPorts) Process() {
	// Dummy
}

func TestConnectInvalidParams(t *testing.T) {
	n := NewGraph()

	n.Add("e1", new(echo))
	n.Add("e2", new(echo))
	n.Add("inv", new(withInvalidPorts))

	cases := []struct {
		scenario string
		err      error
		msg      string
	}{
		{
			"Invalid receiver proc",
			n.Connect("e1", "Out", "noproc", "In"),
			"connect: getProcPort: process 'noproc' not found",
		},
		{
			"Invalid receiver port",
			n.Connect("e1", "Out", "e2", "NotIn"),
			"connect: getProcPort: process 'e2' does not have a valid port 'NotIn'",
		},
		{
			"Invalid sender proc",
			n.Connect("noproc", "Out", "e2", "In"),
			"connect: getProcPort: process 'noproc' not found",
		},
		{
			"Invalid sender port",
			n.Connect("e1", "NotOut", "e2", "In"),
			"connect: getProcPort: process 'e1' does not have a valid port 'NotOut'",
		},
		{
			"Sending to output",
			n.Connect("e1", "Out", "e2", "Out"),
			"connect 'e2.Out': channel does not support direction <-chan",
		},
		{
			"Sending from input",
			n.Connect("e1", "In", "e2", "In"),
			"connect 'e1.In': channel does not support direction chan<-",
		},
		{
			"Connecting to non-chan",
			n.Connect("e1", "Out", "inv", "NotChan"),
			"connect 'inv.NotChan': not a channel",
		},
	}

	for _, item := range cases {
		c := item
		t.Run(c.scenario, func(t *testing.T) {
			t.Parallel()
			if c.err == nil {
				t.Fail()
			} else if c.msg != c.err.Error() {
				t.Error(c.err)
			}
		})
	}
}

func TestSubgraphSender(t *testing.T) {
	sub, err := newDoubleEcho()
	if err != nil {
		t.Error(err)
		return
	}

	n := NewGraph()
	if err := n.Add("sub", sub); err != nil {
		t.Error(err)
		return
	}
	n.Add("e3", new(echo))

	if err := n.Connect("sub", "Out", "e3", "In"); err != nil {
		t.Error(err)
		return
	}

	n.MapInPort("In", "sub", "In")
	n.MapOutPort("Out", "e3", "Out")

	testGraphWithNumberSequence(n, t)
}

func TestSubgraphReceiver(t *testing.T) {
	sub, err := newDoubleEcho()
	if err != nil {
		t.Error(err)
		return
	}

	n := NewGraph()
	if err := n.Add("sub", sub); err != nil {
		t.Error(err)
		return
	}
	n.Add("e3", new(echo))

	if err := n.Connect("e3", "Out", "sub", "In"); err != nil {
		t.Error(err)
		return
	}

	n.MapInPort("In", "e3", "In")
	n.MapOutPort("Out", "sub", "Out")

	testGraphWithNumberSequence(n, t)
}

func newFanOutFanIn() (*Graph, error) {
	n := NewGraph()

	components := map[string]interface{}{
		"e1": new(echo),
		"d1": new(doubler),
		"d2": new(doubler),
		"d3": new(doubler),
		"e2": new(echo),
	}

	for name, c := range components {
		if err := n.Add(name, c); err != nil {
			return nil, err
		}
	}

	connections := []struct{ sn, sp, rn, rp string }{
		{"e1", "Out", "d1", "In"},
		{"e1", "Out", "d2", "In"},
		{"e1", "Out", "d3", "In"},
		{"d1", "Out", "e2", "In"},
		{"d2", "Out", "e2", "In"},
		{"d3", "Out", "e2", "In"},
	}

	for _, c := range connections {
		if err := n.Connect(c.sn, c.sp, c.rn, c.rp); err != nil {
			return nil, err
		}
	}

	n.MapInPort("In", "e1", "In")
	n.MapOutPort("Out", "e2", "Out")

	return n, nil
}

func TestFanOutFanIn(t *testing.T) {
	inData := []int{1, 2, 3, 4, 5, 6, 7, 8}
	outData := []int{2, 4, 6, 8, 10, 12, 14, 16}

	n, err := newFanOutFanIn()
	if err != nil {
		t.Error(err)
		return
	}

	in := make(chan int)
	out := make(chan int)
	n.SetInPort("In", in)
	n.SetOutPort("Out", out)

	wait := Run(n)

	go func() {
		for _, n := range inData {
			in <- n
		}
		close(in)
	}()

	i := 0
	for actual := range out {
		found := false
		for j := 0; j < len(outData); j++ {
			if outData[j] == actual {
				found = true
				outData = append(outData[:j], outData[j+1:]...)
			}
		}
		if !found {
			t.Errorf("%d not found in expected data", actual)
		}
		i++
	}

	if i != len(inData) {
		t.Errorf("Output count missmatch: %d != %d", i, len(inData))
	}

	<-wait
}

func newMapPorts() (*Graph, error) {
	n := NewGraph()

	components := map[string]interface{}{
		"e1":  new(echo),
		"e3":  new(echo),
		"e11": new(echo),
		"e22": new(echo),
		"r":   new(router),
	}

	for name, c := range components {
		if err := n.Add(name, c); err != nil {
			return nil, err
		}
	}

	connections := []struct{ sn, sp, rn, rp string }{
		{"e1", "Out", "r", "In[e1]"},
		{"r", "Out[e3]", "e3", "In"},
		{"r", "Out[e2]", "e22", "In"},
		{"r", "Out[e1]", "e11", "In"},
	}

	for _, c := range connections {
		if err := n.Connect(c.sn, c.sp, c.rn, c.rp); err != nil {
			return nil, err
		}
	}

	iips := []struct {
		proc, port string
		v          int
	}{
		{"e1", "In", 1},
		{"r", "In[e3]", 3},
	}

	for _, p := range iips {
		if err := n.AddIIP(p.proc, p.port, p.v); err != nil {
			return nil, err
		}
	}

	n.MapInPort("I2", "r", "In[e2]")

	outPorts := []struct{ pn, pp, name string }{
		{"e11", "Out", "O1"},
		{"e22", "Out", "O2"},
		{"e3", "Out", "O3"},
	}

	for _, p := range outPorts {
		n.MapOutPort(p.name, p.pn, p.pp)
	}

	return n, nil
}

func TestMapPorts(t *testing.T) {
	n, err := newMapPorts()
	if err != nil {
		t.Error(err)
		return
	}

	i2 := make(chan int, 1)
	o1 := make(chan int)
	o2 := make(chan int)
	o3 := make(chan int)
	if err := n.SetInPort("I2", i2); err != nil {
		t.Error(err)
		return
	}
	n.SetOutPort("O1", o1)
	n.SetOutPort("O2", o2)
	n.SetOutPort("O3", o3)

	wait := Run(n)

	i2 <- 2
	close(i2)
	v1 := <-o1
	v2 := <-o2
	v3 := <-o3

	expected := []int{1, 2, 3}
	actual := []int{v1, v2, v3}

	for i, v := range actual {
		if v != expected[i] {
			t.Errorf("Expected %d, got %d", expected[i], v)
		}
	}

	<-wait
}

func newArrayPorts() (*Graph, error) {
	n := NewGraph()

	components := map[string]interface{}{
		"e0":  new(echo),
		"e2":  new(echo),
		"e00": new(echo),
		"e11": new(echo),
		"r":   new(irouter),
	}

	for name, c := range components {
		if err := n.Add(name, c); err != nil {
			return nil, err
		}
	}

	connections := []struct{ sn, sp, rn, rp string }{
		{"e0", "Out", "r", "In[0]"},
		{"r", "Out[2]", "e2", "In"},
		{"r", "Out[1]", "e11", "In"},
		{"r", "Out[0]", "e00", "In"},
	}

	for _, c := range connections {
		if err := n.Connect(c.sn, c.sp, c.rn, c.rp); err != nil {
			return nil, err
		}
	}

	iips := []struct {
		proc, port string
		v          int
	}{
		{"e0", "In", 1},
		{"r", "In[2]", 3},
	}

	for _, p := range iips {
		if err := n.AddIIP(p.proc, p.port, p.v); err != nil {
			return nil, err
		}
	}

	n.MapInPort("I1", "r", "In[1]")

	outPorts := []struct{ pn, pp, name string }{
		{"e00", "Out", "O0"},
		{"e11", "Out", "O1"},
		{"e2", "Out", "O2"},
	}

	for _, p := range outPorts {
		n.MapOutPort(p.name, p.pn, p.pp)
	}

	return n, nil
}

func TestArrayPorts(t *testing.T) {
	n, err := newArrayPorts()
	if err != nil {
		t.Error(err)
		return
	}

	i1 := make(chan int, 1)
	o0 := make(chan int)
	o1 := make(chan int)
	o2 := make(chan int)
	if err := n.SetInPort("I1", i1); err != nil {
		t.Error(err)
		return
	}
	n.SetOutPort("O0", o0)
	n.SetOutPort("O1", o1)
	n.SetOutPort("O2", o2)

	wait := Run(n)

	i1 <- 2
	close(i1)
	v0 := <-o0
	v1 := <-o1
	v2 := <-o2

	expected := []int{1, 2, 3}
	actual := []int{v0, v1, v2}

	for i, v := range actual {
		if v != expected[i] {
			t.Errorf("Expected %d, got %d", expected[i], v)
		}
	}

	<-wait
}
<commit_msg>Add array/map graph outport to the test<commit_after>package goflow

import (
	"testing"
)

type withInvalidPorts struct {
	NotChan int
	Chan    <-chan int
}

func (c *withInvalidPorts) Process() {
	// Dummy
}

func TestConnectInvalidParams(t *testing.T) {
	n := NewGraph()

	n.Add("e1", new(echo))
	n.Add("e2", new(echo))
	n.Add("inv", new(withInvalidPorts))

	cases := []struct {
		scenario string
		err      error
		msg      string
	}{
		{
			"Invalid receiver proc",
			n.Connect("e1", "Out", "noproc", "In"),
			"connect: getProcPort: process 'noproc' not found",
		},
		{
			"Invalid receiver port",
			n.Connect("e1", "Out", "e2", "NotIn"),
			"connect: getProcPort: process 'e2' does not have a valid port 'NotIn'",
		},
		{
			"Invalid sender proc",
			n.Connect("noproc", "Out", "e2", "In"),
			"connect: getProcPort: process 'noproc' not found",
		},
		{
			"Invalid sender port",
			n.Connect("e1", "NotOut", "e2", "In"),
			"connect: getProcPort: process 'e1' does not have a valid port 'NotOut'",
		},
		{
			"Sending to output",
			n.Connect("e1", "Out", "e2", "Out"),
			"connect 'e2.Out': channel does not support direction <-chan",
		},
		{
			"Sending from input",
			n.Connect("e1", "In", "e2", "In"),
			"connect 'e1.In': channel does not support direction chan<-",
		},
		{
			"Connecting to non-chan",
			n.Connect("e1", "Out", "inv", "NotChan"),
			"connect 'inv.NotChan': not a channel",
		},
	}

	for _, item := range cases {
		c := item
		t.Run(c.scenario, func(t *testing.T) {
			t.Parallel()
			if c.err == nil {
				t.Fail()
			} else if c.msg != c.err.Error() {
				t.Error(c.err)
			}
		})
	}
}

func TestSubgraphSender(t *testing.T) {
	sub, err := newDoubleEcho()
	if err != nil {
		t.Error(err)
		return
	}

	n := NewGraph()
	if err := n.Add("sub", sub); err != nil {
		t.Error(err)
		return
	}
	n.Add("e3", new(echo))

	if err := n.Connect("sub", "Out", "e3", "In"); err != nil {
		t.Error(err)
		return
	}

	n.MapInPort("In", "sub", "In")
	n.MapOutPort("Out", "e3", "Out")

	testGraphWithNumberSequence(n, t)
}

func TestSubgraphReceiver(t *testing.T) {
	sub, err := newDoubleEcho()
	if err != nil {
		t.Error(err)
		return
	}

	n := NewGraph()
	if err := n.Add("sub", sub); err != nil {
		t.Error(err)
		return
	}
	n.Add("e3", new(echo))

	if err := n.Connect("e3", "Out", "sub", "In"); err != nil {
		t.Error(err)
		return
	}

	n.MapInPort("In", "e3", "In")
	n.MapOutPort("Out", "sub", "Out")

	testGraphWithNumberSequence(n, t)
}

func newFanOutFanIn() (*Graph, error) {
	n := NewGraph()

	components := map[string]interface{}{
		"e1": new(echo),
		"d1": new(doubler),
		"d2": new(doubler),
		"d3": new(doubler),
		"e2": new(echo),
	}

	for name, c := range components {
		if err := n.Add(name, c); err != nil {
			return nil, err
		}
	}

	connections := []struct{ sn, sp, rn, rp string }{
		{"e1", "Out", "d1", "In"},
		{"e1", "Out", "d2", "In"},
		{"e1", "Out", "d3", "In"},
		{"d1", "Out", "e2", "In"},
		{"d2", "Out", "e2", "In"},
		{"d3", "Out", "e2", "In"},
	}

	for _, c := range connections {
		if err := n.Connect(c.sn, c.sp, c.rn, c.rp); err != nil {
			return nil, err
		}
	}

	n.MapInPort("In", "e1", "In")
	n.MapOutPort("Out", "e2", "Out")

	return n, nil
}

func TestFanOutFanIn(t *testing.T) {
	inData := []int{1, 2, 3, 4, 5, 6, 7, 8}
	outData := []int{2, 4, 6, 8, 10, 12, 14, 16}

	n, err := newFanOutFanIn()
	if err != nil {
		t.Error(err)
		return
	}

	in := make(chan int)
	out := make(chan int)
	n.SetInPort("In", in)
	n.SetOutPort("Out", out)

	wait := Run(n)

	go func() {
		for _, n := range inData {
			in <- n
		}
		close(in)
	}()

	i := 0
	for actual := range out {
		found := false
		for j := 0; j < len(outData); j++ {
			if outData[j] == actual {
				found = true
				outData = append(outData[:j], outData[j+1:]...)
			}
		}
		if !found {
			t.Errorf("%d not found in expected data", actual)
		}
		i++
	}

	if i != len(inData) {
		t.Errorf("Output count missmatch: %d != %d", i, len(inData))
	}

	<-wait
}

func newMapPorts() (*Graph, error) {
	n := NewGraph()

	components := map[string]interface{}{
		"e1":  new(echo),
		"e11": new(echo),
		"e22": new(echo),
		"r":   new(router),
	}

	for name, c := range components {
		if err := n.Add(name, c); err != nil {
			return nil, err
		}
	}

	connections := []struct{ sn, sp, rn, rp string }{
		{"e1", "Out", "r", "In[e1]"},
		{"r", "Out[e2]", "e22", "In"},
		{"r", "Out[e1]", "e11", "In"},
	}

	for _, c := range connections {
		if err := n.Connect(c.sn, c.sp, c.rn, c.rp); err != nil {
			return nil, err
		}
	}

	iips := []struct {
		proc, port string
		v          int
	}{
		{"e1", "In", 1},
		{"r", "In[e3]", 3},
	}

	for _, p := range iips {
		if err := n.AddIIP(p.proc, p.port, p.v); err != nil {
			return nil, err
		}
	}

	n.MapInPort("I2", "r", "In[e2]")

	outPorts := []struct{ pn, pp, name string }{
		{"e11", "Out", "O1"},
		{"e22", "Out", "O2"},
		{"r", "Out[e3]", "O3"},
	}

	for _, p := range outPorts {
		n.MapOutPort(p.name, p.pn, p.pp)
	}

	return n, nil
}

func TestMapPorts(t *testing.T) {
	n, err := newMapPorts()
	if err != nil {
		t.Error(err)
		return
	}

	i2 := make(chan int, 1)
	o1 := make(chan int)
	o2 := make(chan int)
	o3 := make(chan int)
	if err := n.SetInPort("I2", i2); err != nil {
		t.Error(err)
		return
	}
	n.SetOutPort("O1", o1)
	n.SetOutPort("O2", o2)
	if err := n.SetOutPort("O3", o3); err != nil {
		t.Error(err)
		return
	}

	wait := Run(n)

	i2 <- 2
	close(i2)
	v1 := <-o1
	v2 := <-o2
	v3 := <-o3

	expected := []int{1, 2, 3}
	actual := []int{v1, v2, v3}

	for i, v := range actual {
		if v != expected[i] {
			t.Errorf("Expected %d, got %d", expected[i], v)
		}
	}

	<-wait
}

func newArrayPorts() (*Graph, error) {
	n := NewGraph()

	components := map[string]interface{}{
		"e0":  new(echo),
		"e00": new(echo),
		"e11": new(echo),
		"r":   new(irouter),
	}

	for name, c := range components {
		if err := n.Add(name, c); err != nil {
			return nil, err
		}
	}

	connections := []struct{ sn, sp, rn, rp string }{
		{"e0", "Out", "r", "In[0]"},
		{"r", "Out[1]", "e11", "In"},
		{"r", "Out[0]", "e00", "In"},
	}

	for _, c := range connections {
		if err := n.Connect(c.sn, c.sp, c.rn, c.rp); err != nil {
			return nil, err
		}
	}

	iips := []struct {
		proc, port string
		v          int
	}{
		{"e0", "In", 1},
		{"r", "In[2]", 3},
	}

	for _, p := range iips {
		if err := n.AddIIP(p.proc, p.port, p.v); err != nil {
			return nil, err
		}
	}

	n.MapInPort("I1", "r", "In[1]")

	outPorts := []struct{ pn, pp, name string }{
		{"e00", "Out", "O0"},
		{"e11", "Out", "O1"},
		{"r", "Out[2]", "O2"},
	}

	for _, p := range outPorts {
		n.MapOutPort(p.name, p.pn, p.pp)
	}

	return n, nil
}

func TestArrayPorts(t *testing.T) {
	n, err := newArrayPorts()
	if err != nil {
		t.Error(err)
		return
	}

	i1 := make(chan int, 1)
	o0 := make(chan int)
	o1 := make(chan int)
	o2 := make(chan int)
	if err := n.SetInPort("I1", i1); err != nil {
		t.Error(err)
		return
	}
	n.SetOutPort("O0", o0)
	n.SetOutPort("O1", o1)
	if err := n.SetOutPort("O2", o2); err != nil {
		t.Error(err)
		return
	}

	wait := Run(n)

	i1 <- 2
	close(i1)
	v0 := <-o0
	v1 := <-o1
	v2 := <-o2

	expected := []int{1, 2, 3}
	actual := []int{v0, v1, v2}

	for i, v := range actual {
		if v != expected[i] {
			t.Errorf("Expected %d, got %d", expected[i], v)
		}
	}

	<-wait
}
<|endoftext|>


<commit_before>//
// Copyright 2021, Kordian Bruck
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package gitlab

import (
	"fmt"
	"time"
)

// PackagesService handles communication with the packages related methods
// of the GitLab API.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/packages.html
type PackagesService struct {
	client *Client
}

// Package represents a GitLab single package
//
// GitLab API docs: https://docs.gitlab.com/ee/api/packages.html
type Package struct {
	ID        int        `json:"id"`
	Name      string     `json:"name"`
	Version   string     `json:"version"`
	Type      string     `json:"package_type"`
	CreatedAt *time.Time `json:"created_at"`
}

func (s Package) String() string {
	return Stringify(s)
}

// PackageFile represents one file contained within a package
//
// GitLab API docs: https://docs.gitlab.com/ee/api/packages.html
type PackageFile struct {
	ID        int        `json:"id"`
	PackageID int        `json:"package_id"`
	CreatedAt *time.Time `json:"created_at"`
	FileName  string     `json:"file_name"`
	Size      int        `json:"size"`
	MD5       string     `json:"file_md5"`
	SHA1      string     `json:"file_sha1"`
}

func (s PackageFile) String() string {
	return Stringify(s)
}

// ListProjectPackagesOptions are the parameters available in a ListProjectPackages() Operation
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#within-a-project
type ListProjectPackagesOptions struct {
	ListOptions
	OrderBy            string `url:"order_by,omitempty" json:"order_by,omitempty"`
	Sort               string `url:"sort,omitempty" json:"sort,omitempty"`
	Type               string `url:"package_type,omitempty" json:"package_type,omitempty"`
	Name               string `url:"package_name,omitempty" json:"package_name,omitempty"`
	IncludeVersionless bool   `url:"include_versionless,omitempty" json:"include_versionless,omitempty"`
}

// ListProjectPackages gets a list of packages in a project.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#within-a-project
func (s *PackagesService) ListProjectPackages(pid interface{}, opt *ListProjectPackagesOptions, options ...RequestOptionFunc) ([]*Package, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/packages", pathEscape(project))

	req, err := s.client.NewRequest("GET", u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	var packages []*Package
	resp, err := s.client.Do(req, &packages)
	if err != nil {
		return nil, resp, err
	}

	return packages, resp, err
}

// DeleteRegistryRepository deletes a repository in a registry.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#delete-a-project-package
func (s *PackagesService) DeleteProjectPackage(pid interface{}, packageID int, options ...RequestOptionFunc) (*Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, err
	}
	u := fmt.Sprintf("projects/%s/packages/%d", pathEscape(project), packageID)

	req, err := s.client.NewRequest("DELETE", u, nil, options)
	if err != nil {
		return nil, err
	}

	return s.client.Do(req, nil)
}

// ListPackageFilesOptions represents the available
// ListPackageFiles() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#list-package-files
type ListPackageFilesOptions ListOptions

// ListPackageFiles gets a list of files that are within a package
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#list-package-files
func (s *PackagesService) ListPackageFiles(pid interface{}, packageID int, opt *ListPackageFilesOptions, options ...RequestOptionFunc) ([]*PackageFile, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/packages/%d/package_files",
		pathEscape(project),
		packageID,
	)

	req, err := s.client.NewRequest("GET", u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	var tags []*PackageFile
	resp, err := s.client.Do(req, &tags)
	if err != nil {
		return nil, resp, err
	}

	return tags, resp, err
}
<commit_msg>Update packages.go<commit_after>//
// Copyright 2021, Kordian Bruck
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package gitlab

import (
	"fmt"
	"time"
)

// PackagesService handles communication with the packages related methods
// of the GitLab API.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/packages.html
type PackagesService struct {
	client *Client
}

// Package represents a GitLab single package
//
// GitLab API docs: https://docs.gitlab.com/ee/api/packages.html
type Package struct {
	ID        int        `json:"id"`
	Name      string     `json:"name"`
	Version   string     `json:"version"`
	PackageType      string     `json:"package_type"`
	CreatedAt *time.Time `json:"created_at"`
}

func (s Package) String() string {
	return Stringify(s)
}

// PackageFile represents one file contained within a package
//
// GitLab API docs: https://docs.gitlab.com/ee/api/packages.html
type PackageFile struct {
	ID        int        `json:"id"`
	PackageID int        `json:"package_id"`
	CreatedAt *time.Time `json:"created_at"`
	FileName  string     `json:"file_name"`
	Size      int        `json:"size"`
	MD5       string     `json:"file_md5"`
	SHA1      string     `json:"file_sha1"`
}

func (s PackageFile) String() string {
	return Stringify(s)
}

// ListProjectPackagesOptions are the parameters available in a ListProjectPackages() Operation
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#within-a-project
type ListProjectPackagesOptions struct {
	ListOptions
	OrderBy            string `url:"order_by,omitempty" json:"order_by,omitempty"`
	Sort               string `url:"sort,omitempty" json:"sort,omitempty"`
	Type               string `url:"package_type,omitempty" json:"package_type,omitempty"`
	Name               string `url:"package_name,omitempty" json:"package_name,omitempty"`
	IncludeVersionless bool   `url:"include_versionless,omitempty" json:"include_versionless,omitempty"`
}

// ListProjectPackages gets a list of packages in a project.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#within-a-project
func (s *PackagesService) ListProjectPackages(pid interface{}, opt *ListProjectPackagesOptions, options ...RequestOptionFunc) ([]*Package, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/packages", pathEscape(project))

	req, err := s.client.NewRequest("GET", u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	var packages []*Package
	resp, err := s.client.Do(req, &packages)
	if err != nil {
		return nil, resp, err
	}

	return packages, resp, err
}

// DeleteRegistryRepository deletes a repository in a registry.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#delete-a-project-package
func (s *PackagesService) DeleteProjectPackage(pid interface{}, packageID int, options ...RequestOptionFunc) (*Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, err
	}
	u := fmt.Sprintf("projects/%s/packages/%d", pathEscape(project), packageID)

	req, err := s.client.NewRequest("DELETE", u, nil, options)
	if err != nil {
		return nil, err
	}

	return s.client.Do(req, nil)
}

// ListPackageFilesOptions represents the available
// ListPackageFiles() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#list-package-files
type ListPackageFilesOptions ListOptions

// ListPackageFiles gets a list of files that are within a package
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/packages.html#list-package-files
func (s *PackagesService) ListPackageFiles(pid interface{}, packageID int, opt *ListPackageFilesOptions, options ...RequestOptionFunc) ([]*PackageFile, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/packages/%d/package_files",
		pathEscape(project),
		packageID,
	)

	req, err := s.client.NewRequest("GET", u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	var tags []*PackageFile
	resp, err := s.client.Do(req, &tags)
	if err != nil {
		return nil, resp, err
	}

	return tags, resp, err
}
<|endoftext|>


<commit_before>package filer

import (
	"bytes"
	"fmt"
	"io"
	"math"
	"strings"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
)

func StreamContent(masterClient wdclient.HasLookupFileIdFunction, w io.Writer, chunks []*filer_pb.FileChunk, offset int64, size int64) error {

	glog.V(9).Infof("start to stream content for chunks: %+v\n", chunks)
	chunkViews := ViewFromChunks(masterClient.GetLookupFileIdFunction(), chunks, offset, size)

	fileId2Url := make(map[string][]string)

	for _, chunkView := range chunkViews {

		urlStrings, err := masterClient.GetLookupFileIdFunction()(chunkView.FileId)
		if err != nil {
			glog.V(1).Infof("operation LookupFileId %s failed, err: %v", chunkView.FileId, err)
			return err
		} else if len(urlStrings) == 0 {
			glog.Errorf("operation LookupFileId %s failed, err: urls not found", chunkView.FileId)
			return fmt.Errorf("operation LookupFileId %s failed, err: urls not found", chunkView.FileId)
		}
		fileId2Url[chunkView.FileId] = urlStrings
	}

	for _, chunkView := range chunkViews {

		urlStrings := fileId2Url[chunkView.FileId]
		start := time.Now()
		data, err := retriedFetchChunkData(urlStrings, chunkView.CipherKey, chunkView.IsGzipped, chunkView.IsFullChunk(), chunkView.Offset, int(chunkView.Size))
		stats.FilerRequestHistogram.WithLabelValues("chunkDownload").Observe(time.Since(start).Seconds())
		if err != nil {
			stats.FilerRequestCounter.WithLabelValues("chunkDownloadError").Inc()
			return fmt.Errorf("read chunk: %v", err)
		}

		_, err = w.Write(data)
		if err != nil {
			stats.FilerRequestCounter.WithLabelValues("chunkDownloadedError").Inc()
			return fmt.Errorf("write chunk: %v", err)
		}
		stats.FilerRequestCounter.WithLabelValues("chunkDownload").Inc()
	}

	return nil

}

// ----------------  ReadAllReader ----------------------------------

func ReadAll(masterClient *wdclient.MasterClient, chunks []*filer_pb.FileChunk) ([]byte, error) {

	buffer := bytes.Buffer{}

	lookupFileIdFn := func(fileId string) (targetUrls []string, err error) {
		return masterClient.LookupFileId(fileId)
	}

	chunkViews := ViewFromChunks(lookupFileIdFn, chunks, 0, math.MaxInt64)

	for _, chunkView := range chunkViews {
		urlStrings, err := lookupFileIdFn(chunkView.FileId)
		if err != nil {
			glog.V(1).Infof("operation LookupFileId %s failed, err: %v", chunkView.FileId, err)
			return nil, err
		}

		data, err := retriedFetchChunkData(urlStrings, chunkView.CipherKey, chunkView.IsGzipped, chunkView.IsFullChunk(), chunkView.Offset, int(chunkView.Size))
		if err != nil {
			return nil, err
		}
		buffer.Write(data)
	}
	return buffer.Bytes(), nil
}

// ----------------  ChunkStreamReader ----------------------------------
type ChunkStreamReader struct {
	chunkViews   []*ChunkView
	logicOffset  int64
	buffer       []byte
	bufferOffset int64
	bufferPos    int
	chunkIndex   int
	lookupFileId wdclient.LookupFileIdFunctionType
}

var _ = io.ReadSeeker(&ChunkStreamReader{})

func NewChunkStreamReaderFromFiler(masterClient *wdclient.MasterClient, chunks []*filer_pb.FileChunk) *ChunkStreamReader {

	lookupFileIdFn := func(fileId string) (targetUrl []string, err error) {
		return masterClient.LookupFileId(fileId)
	}

	chunkViews := ViewFromChunks(lookupFileIdFn, chunks, 0, math.MaxInt64)

	return &ChunkStreamReader{
		chunkViews:   chunkViews,
		lookupFileId: lookupFileIdFn,
	}
}

func NewChunkStreamReader(filerClient filer_pb.FilerClient, chunks []*filer_pb.FileChunk) *ChunkStreamReader {

	lookupFileIdFn := LookupFn(filerClient)

	chunkViews := ViewFromChunks(lookupFileIdFn, chunks, 0, math.MaxInt64)

	return &ChunkStreamReader{
		chunkViews:   chunkViews,
		lookupFileId: lookupFileIdFn,
	}
}

func (c *ChunkStreamReader) Read(p []byte) (n int, err error) {
	for n < len(p) {
		if c.isBufferEmpty() {
			if c.chunkIndex >= len(c.chunkViews) {
				return n, io.EOF
			}
			chunkView := c.chunkViews[c.chunkIndex]
			c.fetchChunkToBuffer(chunkView)
			c.chunkIndex++
		}
		t := copy(p[n:], c.buffer[c.bufferPos:])
		c.bufferPos += t
		n += t
	}
	return
}

func (c *ChunkStreamReader) isBufferEmpty() bool {
	return len(c.buffer) <= c.bufferPos
}

func (c *ChunkStreamReader) Seek(offset int64, whence int) (int64, error) {

	var totalSize int64
	for _, chunk := range c.chunkViews {
		totalSize += int64(chunk.Size)
	}

	var err error
	switch whence {
	case io.SeekStart:
	case io.SeekCurrent:
		offset += c.bufferOffset + int64(c.bufferPos)
	case io.SeekEnd:
		offset = totalSize + offset
	}
	if offset > totalSize {
		err = io.ErrUnexpectedEOF
	}

	for i, chunk := range c.chunkViews {
		if chunk.LogicOffset <= offset && offset < chunk.LogicOffset+int64(chunk.Size) {
			if c.isBufferEmpty() || c.bufferOffset != chunk.LogicOffset {
				c.fetchChunkToBuffer(chunk)
				c.chunkIndex = i + 1
				break
			}
		}
	}
	c.bufferPos = int(offset - c.bufferOffset)

	return offset, err

}

func (c *ChunkStreamReader) fetchChunkToBuffer(chunkView *ChunkView) error {
	urlStrings, err := c.lookupFileId(chunkView.FileId)
	if err != nil {
		glog.V(1).Infof("operation LookupFileId %s failed, err: %v", chunkView.FileId, err)
		return err
	}
	var buffer bytes.Buffer
	var shouldRetry bool
	for _, urlString := range urlStrings {
		shouldRetry, err = util.ReadUrlAsStream(urlString, chunkView.CipherKey, chunkView.IsGzipped, chunkView.IsFullChunk(), chunkView.Offset, int(chunkView.Size), func(data []byte) {
			buffer.Write(data)
		})
		if !shouldRetry {
			break
		}
		if err != nil {
			glog.V(1).Infof("read %s failed, err: %v", chunkView.FileId, err)
			buffer.Reset()
		} else {
			break
		}
	}
	if err != nil {
		return err
	}
	c.buffer = buffer.Bytes()
	c.bufferPos = 0
	c.bufferOffset = chunkView.LogicOffset

	// glog.V(0).Infof("read %s [%d,%d)", chunkView.FileId, chunkView.LogicOffset, chunkView.LogicOffset+int64(chunkView.Size))

	return nil
}

func (c *ChunkStreamReader) Close() {
	// TODO try to release and reuse buffer
}

func VolumeId(fileId string) string {
	lastCommaIndex := strings.LastIndex(fileId, ",")
	if lastCommaIndex > 0 {
		return fileId[:lastCommaIndex]
	}
	return fileId
}
<commit_msg>refactor<commit_after>package filer

import (
	"bytes"
	"fmt"
	"io"
	"math"
	"strings"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
)

func StreamContent(masterClient wdclient.HasLookupFileIdFunction, w io.Writer, chunks []*filer_pb.FileChunk, offset int64, size int64) error {

	glog.V(9).Infof("start to stream content for chunks: %+v\n", chunks)
	chunkViews := ViewFromChunks(masterClient.GetLookupFileIdFunction(), chunks, offset, size)

	fileId2Url := make(map[string][]string)

	for _, chunkView := range chunkViews {

		urlStrings, err := masterClient.GetLookupFileIdFunction()(chunkView.FileId)
		if err != nil {
			glog.V(1).Infof("operation LookupFileId %s failed, err: %v", chunkView.FileId, err)
			return err
		} else if len(urlStrings) == 0 {
			glog.Errorf("operation LookupFileId %s failed, err: urls not found", chunkView.FileId)
			return fmt.Errorf("operation LookupFileId %s failed, err: urls not found", chunkView.FileId)
		}
		fileId2Url[chunkView.FileId] = urlStrings
	}

	for _, chunkView := range chunkViews {

		urlStrings := fileId2Url[chunkView.FileId]
		start := time.Now()
		data, err := retriedFetchChunkData(urlStrings, chunkView.CipherKey, chunkView.IsGzipped, chunkView.IsFullChunk(), chunkView.Offset, int(chunkView.Size))
		stats.FilerRequestHistogram.WithLabelValues("chunkDownload").Observe(time.Since(start).Seconds())
		if err != nil {
			stats.FilerRequestCounter.WithLabelValues("chunkDownloadError").Inc()
			return fmt.Errorf("read chunk: %v", err)
		}

		_, err = w.Write(data)
		if err != nil {
			stats.FilerRequestCounter.WithLabelValues("chunkDownloadedError").Inc()
			return fmt.Errorf("write chunk: %v", err)
		}
		stats.FilerRequestCounter.WithLabelValues("chunkDownload").Inc()
	}

	return nil

}

// ----------------  ReadAllReader ----------------------------------

func ReadAll(masterClient *wdclient.MasterClient, chunks []*filer_pb.FileChunk) ([]byte, error) {

	buffer := bytes.Buffer{}

	lookupFileIdFn := func(fileId string) (targetUrls []string, err error) {
		return masterClient.LookupFileId(fileId)
	}

	chunkViews := ViewFromChunks(lookupFileIdFn, chunks, 0, math.MaxInt64)

	for _, chunkView := range chunkViews {
		urlStrings, err := lookupFileIdFn(chunkView.FileId)
		if err != nil {
			glog.V(1).Infof("operation LookupFileId %s failed, err: %v", chunkView.FileId, err)
			return nil, err
		}

		data, err := retriedFetchChunkData(urlStrings, chunkView.CipherKey, chunkView.IsGzipped, chunkView.IsFullChunk(), chunkView.Offset, int(chunkView.Size))
		if err != nil {
			return nil, err
		}
		buffer.Write(data)
	}
	return buffer.Bytes(), nil
}

// ----------------  ChunkStreamReader ----------------------------------
type ChunkStreamReader struct {
	chunkViews   []*ChunkView
	totalSize    int64
	logicOffset  int64
	buffer       []byte
	bufferOffset int64
	bufferPos    int
	chunkIndex   int
	lookupFileId wdclient.LookupFileIdFunctionType
}

var _ = io.ReadSeeker(&ChunkStreamReader{})

func NewChunkStreamReaderFromFiler(masterClient *wdclient.MasterClient, chunks []*filer_pb.FileChunk) *ChunkStreamReader {

	lookupFileIdFn := func(fileId string) (targetUrl []string, err error) {
		return masterClient.LookupFileId(fileId)
	}

	chunkViews := ViewFromChunks(lookupFileIdFn, chunks, 0, math.MaxInt64)

	var totalSize int64
	for _, chunk := range chunkViews {
		totalSize += int64(chunk.Size)
	}

	return &ChunkStreamReader{
		chunkViews:   chunkViews,
		lookupFileId: lookupFileIdFn,
		totalSize:    totalSize,
	}
}

func NewChunkStreamReader(filerClient filer_pb.FilerClient, chunks []*filer_pb.FileChunk) *ChunkStreamReader {

	lookupFileIdFn := LookupFn(filerClient)

	chunkViews := ViewFromChunks(lookupFileIdFn, chunks, 0, math.MaxInt64)

	var totalSize int64
	for _, chunk := range chunkViews {
		totalSize += int64(chunk.Size)
	}

	return &ChunkStreamReader{
		chunkViews:   chunkViews,
		lookupFileId: lookupFileIdFn,
		totalSize:    totalSize,
	}
}

func (c *ChunkStreamReader) Read(p []byte) (n int, err error) {
	for n < len(p) {
		if c.isBufferEmpty() {
			if c.chunkIndex >= len(c.chunkViews) {
				return n, io.EOF
			}
			chunkView := c.chunkViews[c.chunkIndex]
			c.fetchChunkToBuffer(chunkView)
			c.chunkIndex++
		}
		t := copy(p[n:], c.buffer[c.bufferPos:])
		c.bufferPos += t
		n += t
	}
	return
}

func (c *ChunkStreamReader) isBufferEmpty() bool {
	return len(c.buffer) <= c.bufferPos
}

func (c *ChunkStreamReader) Seek(offset int64, whence int) (int64, error) {

	var err error
	switch whence {
	case io.SeekStart:
	case io.SeekCurrent:
		offset += c.bufferOffset + int64(c.bufferPos)
	case io.SeekEnd:
		offset = c.totalSize + offset
	}
	if offset > c.totalSize {
		err = io.ErrUnexpectedEOF
	}

	for i, chunk := range c.chunkViews {
		if chunk.LogicOffset <= offset && offset < chunk.LogicOffset+int64(chunk.Size) {
			if c.isBufferEmpty() || c.bufferOffset != chunk.LogicOffset {
				c.fetchChunkToBuffer(chunk)
				c.chunkIndex = i + 1
				break
			}
		}
	}
	c.bufferPos = int(offset - c.bufferOffset)

	return offset, err

}

func (c *ChunkStreamReader) fetchChunkToBuffer(chunkView *ChunkView) error {
	urlStrings, err := c.lookupFileId(chunkView.FileId)
	if err != nil {
		glog.V(1).Infof("operation LookupFileId %s failed, err: %v", chunkView.FileId, err)
		return err
	}
	var buffer bytes.Buffer
	var shouldRetry bool
	for _, urlString := range urlStrings {
		shouldRetry, err = util.ReadUrlAsStream(urlString, chunkView.CipherKey, chunkView.IsGzipped, chunkView.IsFullChunk(), chunkView.Offset, int(chunkView.Size), func(data []byte) {
			buffer.Write(data)
		})
		if !shouldRetry {
			break
		}
		if err != nil {
			glog.V(1).Infof("read %s failed, err: %v", chunkView.FileId, err)
			buffer.Reset()
		} else {
			break
		}
	}
	if err != nil {
		return err
	}
	c.buffer = buffer.Bytes()
	c.bufferPos = 0
	c.bufferOffset = chunkView.LogicOffset

	// glog.V(0).Infof("read %s [%d,%d)", chunkView.FileId, chunkView.LogicOffset, chunkView.LogicOffset+int64(chunkView.Size))

	return nil
}

func (c *ChunkStreamReader) Close() {
	// TODO try to release and reuse buffer
}

func VolumeId(fileId string) string {
	lastCommaIndex := strings.LastIndex(fileId, ",")
	if lastCommaIndex > 0 {
		return fileId[:lastCommaIndex]
	}
	return fileId
}
<|endoftext|>