<commit_before>package inbucket

import (
	"bytes"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/mail"
	"os"
)

const MIME_BUF_BYTES = 1024

type MIMEMessage struct {
	text string
	html string
}

type MIMENode struct {
	Parent      *MIMENode
	FirstChild  *MIMENode
	NextSibling *MIMENode
	Type        string
	Content     *bytes.Buffer
}

func NewMIMENode(parent *MIMENode, contentType string) *MIMENode {
	return &MIMENode{Parent: parent, Type: contentType}
}

func (n *MIMENode) String() string {
	children := ""
	siblings := ""
	if n.FirstChild != nil {
		children = n.FirstChild.String()
	}
	if n.NextSibling != nil {
		siblings = n.NextSibling.String()
	}
	return fmt.Sprintf("[%v %v] %v", n.Type, children, siblings)
}

func ParseMIMEMessage(mailMsg *mail.Message) (*MIMEMessage, error) {
	mimeMsg := new(MIMEMessage)

	// Parse top-level multipart
	ctype := mailMsg.Header.Get("Content-Type")
	mediatype, params, err := mime.ParseMediaType(ctype)
	if err != nil {
		return nil, err
	}
	switch mediatype {
	case "multipart/alternative":
		// Good
	default:
		return nil, fmt.Errorf("Unknown mediatype: %v", mediatype)
	}
	boundary := params["boundary"]
	if boundary == "" {
		return nil, fmt.Errorf("Unable to locate boundary param in Content-Type header")
	}

	// Root Node of our tree
	root := NewMIMENode(nil, mediatype)

	err = parseNodes(root, mailMsg.Body, boundary)
	fmt.Println(root.String())
	return mimeMsg, err
}

func parseNodes(parent *MIMENode, reader io.Reader, boundary string) error {
	var prevSibling *MIMENode

	// Loop over MIME parts
	mr := multipart.NewReader(reader, boundary)
	for {
		part, err := mr.NextPart()
		if err != nil {
			if err == io.EOF {
				// This is a clean end-of-message signal
				break
			}
			return err
		}
		mediatype, params, err := mime.ParseMediaType(part.Header.Get("Content-Type"))
		if err != nil {
			return err
		}

		// Insert ourselves into tree
		node := NewMIMENode(parent, mediatype)
		if prevSibling != nil {
			prevSibling.NextSibling = node
		} else {
			parent.FirstChild = node
		}
		prevSibling = node

		boundary := params["boundary"]
		if boundary != "" {
			// Content is another multipart
			err = parseNodes(node, part, boundary)
			if err != nil {
				return err
			}
		} else {
			// Content is data, allocate a buffer
			node.Content = new(bytes.Buffer)
			_, err = io.Copy(node.Content, part)
			if err != nil {
				return err
			}
			fmt.Printf("\n----\n")
			io.Copy(os.Stdout, node.Content)
			fmt.Printf("\n----\n")
		}
	}

	return nil
}
<commit_msg>More MIME parsing progress!<commit_after>package inbucket

import (
	"bytes"
	"container/list"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/mail"
)

type MIMENodeMatcher func(node *MIMENode) bool

type MIMENode struct {
	Parent      *MIMENode
	FirstChild  *MIMENode
	NextSibling *MIMENode
	Type        string
	Content     []byte
}

type MIMEMessage struct {
	Text string
	Html string
	Root *MIMENode
}

func NewMIMENode(parent *MIMENode, contentType string) *MIMENode {
	return &MIMENode{Parent: parent, Type: contentType}
}

func (n *MIMENode) BreadthFirstSearch(matcher MIMENodeMatcher) *MIMENode {
	q := list.New()
	q.PushBack(n)

	// Push children onto queue and attempt to match in that order
	for q.Len() > 0 {
		e := q.Front()
		n := e.Value.(*MIMENode)
		if matcher(n) {
			return n
		}
		q.Remove(e)
		c := n.FirstChild
		for c != nil {
			q.PushBack(c)
			c = c.NextSibling
		}
	}

	return nil
}

func (n *MIMENode) String() string {
	children := ""
	siblings := ""
	if n.FirstChild != nil {
		children = n.FirstChild.String()
	}
	if n.NextSibling != nil {
		siblings = n.NextSibling.String()
	}
	return fmt.Sprintf("[%v %v] %v", n.Type, children, siblings)
}

func ParseMIMEMessage(mailMsg *mail.Message) (*MIMEMessage, error) {
	mimeMsg := new(MIMEMessage)

	// Parse top-level multipart
	ctype := mailMsg.Header.Get("Content-Type")
	mediatype, params, err := mime.ParseMediaType(ctype)
	if err != nil {
		return nil, err
	}
	switch mediatype {
	case "multipart/alternative":
		// Good
	default:
		return nil, fmt.Errorf("Unknown mediatype: %v", mediatype)
	}
	boundary := params["boundary"]
	if boundary == "" {
		return nil, fmt.Errorf("Unable to locate boundary param in Content-Type header")
	}

	// Root Node of our tree
	root := NewMIMENode(nil, mediatype)

	err = parseNodes(root, mailMsg.Body, boundary)
	fmt.Println(root.String())

	// Locate text body
	match := root.BreadthFirstSearch(func(node *MIMENode) bool {
		return node.Type == "text/plain"
	})
	if match != nil {
		mimeMsg.Text = string(match.Content)
	}

	// Locate HTML body
	match = root.BreadthFirstSearch(func(node *MIMENode) bool {
		return node.Type == "text/html"
	})
	if match != nil {
		mimeMsg.Html = string(match.Content)
	}

	fmt.Println(mimeMsg.String())

	return mimeMsg, err
}

func (m *MIMEMessage) String() string {
	return fmt.Sprintf("----TEXT----\n%v\n----HTML----\n%v\n----END----\n", m.Text, m.Html)
}

func parseNodes(parent *MIMENode, reader io.Reader, boundary string) error {
	var prevSibling *MIMENode

	// Loop over MIME parts
	mr := multipart.NewReader(reader, boundary)
	for {
		part, err := mr.NextPart()
		if err != nil {
			if err == io.EOF {
				// This is a clean end-of-message signal
				break
			}
			return err
		}
		mediatype, params, err := mime.ParseMediaType(part.Header.Get("Content-Type"))
		if err != nil {
			return err
		}

		// Insert ourselves into tree
		node := NewMIMENode(parent, mediatype)
		if prevSibling != nil {
			prevSibling.NextSibling = node
		} else {
			parent.FirstChild = node
		}
		prevSibling = node

		boundary := params["boundary"]
		if boundary != "" {
			// Content is another multipart
			err = parseNodes(node, part, boundary)
			if err != nil {
				return err
			}
		} else {
			// Content is data, allocate a buffer
			buf := new(bytes.Buffer)
			_, err = buf.ReadFrom(part)
			if err != nil {
				return err
			}
			node.Content = buf.Bytes()
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/proxypoke/i3ipc"
	"os"
	"os/exec"
	"strings"
)

func dfsTree(t *i3ipc.I3Node) (names []string, nodes []*i3ipc.I3Node) {
	if t.Layout == "dockarea" {
		return
	}

	if t.Geometry.Width > 0 && t.Geometry.Height > 0 && t.Window > 0 && !t.Focused {
		name := t.Name
		if len(name) == 0 {
			name = fmt.Sprint("(Unnamed) ", t.Window)
		}
		names = append(names, name)
		nodes = append(nodes, t)
	}
	visitChildren := func(children *[]i3ipc.I3Node) {
		for i, _ := range *children {
			newNames, newNodes := dfsTree(&(*children)[i])
			names = append(names, newNames...)
			nodes = append(nodes, newNodes...)
		}
	}
	visitChildren(&t.FloatingNodes)
	visitChildren(&t.Nodes)
	return
}

func checkError(e error) {
	if e != nil {
		fmt.Fprintln(os.Stderr, e)
		os.Exit(0)
	}
}

func main() {
	ipc, e := i3ipc.GetIPCSocket()
	checkError(e)

	tree, e := ipc.GetTree()
	checkError(e)

	var node *i3ipc.I3Node
	names, nodes := dfsTree(&tree)
	switch len(names) {
	case 0:
		os.Exit(0)
	case 1:
		node = nodes[0]
	default:
		cmd := exec.Command("dmenu", os.Args[1:len(os.Args)]...)
		cmd.Stdin = strings.NewReader(strings.Join(names, "\n"))

		out, e := cmd.Output()
		checkError(e)

		name := strings.TrimRight(string(out), "\n")

		for i, s := range names {
			if s == name {
				node = nodes[i]
				break
			}
		}
	}

	if node != nil {
		msg := fmt.Sprint("[con_id=", node.Id, "] focus")
		ipc.Raw(i3ipc.I3Command, msg)
	}
}
<commit_msg>Exit 1 when error happens<commit_after>package main

import (
	"fmt"
	"github.com/proxypoke/i3ipc"
	"os"
	"os/exec"
	"strings"
)

func dfsTree(t *i3ipc.I3Node) (names []string, nodes []*i3ipc.I3Node) {
	if t.Layout == "dockarea" {
		return
	}

	if t.Geometry.Width > 0 && t.Geometry.Height > 0 && t.Window > 0 && !t.Focused {
		name := t.Name
		if len(name) == 0 {
			name = fmt.Sprint("(Unnamed) ", t.Window)
		}
		names = append(names, name)
		nodes = append(nodes, t)
	}
	visitChildren := func(children *[]i3ipc.I3Node) {
		for i, _ := range *children {
			newNames, newNodes := dfsTree(&(*children)[i])
			names = append(names, newNames...)
			nodes = append(nodes, newNodes...)
		}
	}
	visitChildren(&t.FloatingNodes)
	visitChildren(&t.Nodes)
	return
}

func checkError(e error) {
	if e != nil {
		fmt.Fprintln(os.Stderr, e)
		os.Exit(1)
	}
}

func main() {
	ipc, e := i3ipc.GetIPCSocket()
	checkError(e)

	tree, e := ipc.GetTree()
	checkError(e)

	var node *i3ipc.I3Node
	names, nodes := dfsTree(&tree)
	switch len(names) {
	case 0:
		os.Exit(0)
	case 1:
		node = nodes[0]
	default:
		cmd := exec.Command("dmenu", os.Args[1:len(os.Args)]...)
		cmd.Stdin = strings.NewReader(strings.Join(names, "\n"))

		out, e := cmd.Output()
		checkError(e)

		name := strings.TrimRight(string(out), "\n")

		for i, s := range names {
			if s == name {
				node = nodes[i]
				break
			}
		}
	}

	if node != nil {
		msg := fmt.Sprint("[con_id=", node.Id, "] focus")
		ipc.Raw(i3ipc.I3Command, msg)
	}
}
<|endoftext|>


<commit_before>package fakeclock_test

import (
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-golang/clock/fakeclock"
)

var _ = Describe("FakeClock", func() {
	const Δ time.Duration = 10 * time.Millisecond

	var (
		fakeClock   *fakeclock.FakeClock
		initialTime time.Time
	)

	BeforeEach(func() {
		initialTime = time.Date(2014, 1, 1, 3, 0, 30, 0, time.UTC)
		fakeClock = fakeclock.NewFakeClock(initialTime)
	})

	Describe("Now", func() {
		It("returns the current time, w/o race conditions", func() {
			go fakeClock.Increment(time.Minute)
			Eventually(fakeClock.Now).Should(Equal(initialTime.Add(time.Minute)))
		})
	})

	Describe("Sleep", func() {
		It("blocks until the given interval elapses", func() {
			doneSleeping := make(chan struct{})
			go func() {
				fakeClock.Sleep(10 * time.Second)
				close(doneSleeping)
			}()

			Consistently(doneSleeping, Δ).ShouldNot(BeClosed())

			fakeClock.Increment(5 * time.Second)
			Consistently(doneSleeping, Δ).ShouldNot(BeClosed())

			fakeClock.Increment(4 * time.Second)
			Consistently(doneSleeping, Δ).ShouldNot(BeClosed())

			fakeClock.Increment(1 * time.Second)
			Eventually(doneSleeping).Should(BeClosed())
		})
	})

	Describe("WatcherCount", func() {
		Context("when a timer is created", func() {
			It("increments the watcher count", func() {
				fakeClock.NewTimer(time.Second)
				Ω(fakeClock.WatcherCount()).Should(Equal(1))

				fakeClock.NewTimer(2 * time.Second)
				Ω(fakeClock.WatcherCount()).Should(Equal(2))
			})
		})

		Context("when a timer fires", func() {
			It("increments the watcher count", func() {
				fakeClock.NewTimer(time.Second)
				Ω(fakeClock.WatcherCount()).Should(Equal(1))

				fakeClock.Increment(time.Second)
				Ω(fakeClock.WatcherCount()).Should(Equal(0))
			})
		})
	})
})
<commit_msg>Change Ω to Expect [#93246926]<commit_after>package fakeclock_test

import (
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-golang/clock/fakeclock"
)

var _ = Describe("FakeClock", func() {
	const Δ time.Duration = 10 * time.Millisecond

	var (
		fakeClock   *fakeclock.FakeClock
		initialTime time.Time
	)

	BeforeEach(func() {
		initialTime = time.Date(2014, 1, 1, 3, 0, 30, 0, time.UTC)
		fakeClock = fakeclock.NewFakeClock(initialTime)
	})

	Describe("Now", func() {
		It("returns the current time, w/o race conditions", func() {
			go fakeClock.Increment(time.Minute)
			Eventually(fakeClock.Now).Should(Equal(initialTime.Add(time.Minute)))
		})
	})

	Describe("Sleep", func() {
		It("blocks until the given interval elapses", func() {
			doneSleeping := make(chan struct{})
			go func() {
				fakeClock.Sleep(10 * time.Second)
				close(doneSleeping)
			}()

			Consistently(doneSleeping, Δ).ShouldNot(BeClosed())

			fakeClock.Increment(5 * time.Second)
			Consistently(doneSleeping, Δ).ShouldNot(BeClosed())

			fakeClock.Increment(4 * time.Second)
			Consistently(doneSleeping, Δ).ShouldNot(BeClosed())

			fakeClock.Increment(1 * time.Second)
			Eventually(doneSleeping).Should(BeClosed())
		})
	})

	Describe("WatcherCount", func() {
		Context("when a timer is created", func() {
			It("increments the watcher count", func() {
				fakeClock.NewTimer(time.Second)
				Expect(fakeClock.WatcherCount()).To(Equal(1))

				fakeClock.NewTimer(2 * time.Second)
				Expect(fakeClock.WatcherCount()).To(Equal(2))
			})
		})

		Context("when a timer fires", func() {
			It("increments the watcher count", func() {
				fakeClock.NewTimer(time.Second)
				Expect(fakeClock.WatcherCount()).To(Equal(1))

				fakeClock.Increment(time.Second)
				Expect(fakeClock.WatcherCount()).To(Equal(0))
			})
		})
	})
})
<|endoftext|>


<commit_before>package scanner

import (
	"fmt"
	"reflect"
	"strings"
	"sync"
	"testing"
)

var tests = []struct {
	src     string
	results []Token
}{
	{
		src: "export fn\nfn1",
		results: []Token{
			Token{Type: TokenTypeIdent, Text: "export"},
			Token{Type: TokenTypeWhitespace, Column: 6, Text: " "},
			Token{Type: TokenTypeIdent, Column: 7, Text: "fn"},
			Token{Type: TokenTypeWhitespace, Column: 9, Text: "\n"},
			Token{Type: TokenTypeIdent, Line: 1, Text: "fn1"},
			Token{Type: TokenTypeEOF, Line: 1, Column: 3, Text: ""},
		},
	},
	{
		src: "\"foobar\"\n\"foo",
		results: []Token{
			Token{Type: TokenTypeString, Text: `"foobar"`, Value: "foobar"},
			Token{Type: TokenTypeWhitespace, Column: 8, Text: "\n"},
			Token{Type: TokenTypeUnknown, Line: 1, Column: 0, Text: `"foo`, Value: ""},
			Token{Type: TokenTypeEOF, Line: 1, Column: 4, Text: ``},
		},
	},
	{
		src: `"foo\"bar""foo\"bar"`,
		results: []Token{
			Token{Type: TokenTypeString, Text: `"foo\"bar"`, Value: `foo"bar`},
			Token{Type: TokenTypeString, Column: 10, Text: `"foo\"bar"`, Value: `foo"bar`},
			Token{Type: TokenTypeEOF, Line: 0, Column: 20, Text: ``},
		},
	},
	{
		src: "\"foo\\",
		results: []Token{
			Token{Type: TokenTypeUnknown, Text: `"foo\`, Value: ""},
			Token{Type: TokenTypeEOF, Line: 0, Column: 5, Text: ``},
		},
	},
	{
		src: `"foo\nbar"`,
		results: []Token{
			Token{Type: TokenTypeString, Text: `"foo\nbar"`, Value: "foo\\nbar"},
			Token{Type: TokenTypeEOF, Line: 0, Column: 10, Text: ``},
		},
	},
	{
		src: `"foo\\bar"`,
		results: []Token{
			Token{Type: TokenTypeString, Text: `"foo\\bar"`, Value: `foo\bar`},
			Token{Type: TokenTypeEOF, Line: 0, Column: 10, Text: ``},
		},
	},
	{
		src: "\"\\123\\x53\\u2318\"",
		results: []Token{
			Token{Type: TokenTypeString, Text: "\"\\123\\x53\\u2318\"", Value: "SS⌘"},
			Token{Type: TokenTypeEOF, Line: 0, Column: 16, Text: ``},
		},
	},
	/*
		{
			src: `"\xe6\x97\xa5\xe6\x9c\xac\xe8\xaa\x9e"`,
			results: []Token{
				Token{Type: TokenTypeString, Text: `"\xe6\x97\xa5\xe6\x9c\xac\xe8\xaa\x9e"`, Value: "日本語"},
				Token{Type: TokenTypeEOF, Line: 0, Column: 38, Text: ``},
			},
		},
	*/
	{
		src: "`\\n\\000`",
		results: []Token{
			Token{Type: TokenTypeString, Text: "`\\n\\000`", Value: "\\n\\000"},
			Token{Type: TokenTypeEOF, Line: 0, Column: 8, Text: ``},
		},
	},
	{
		src: "\"\n\"",
		results: []Token{
			Token{Type: TokenTypeUnknown, Text: "\"", Value: ""},
			Token{Type: TokenTypeWhitespace, Line: 0, Column: 1, Text: "\n", Value: nil},
			Token{Type: TokenTypeUnknown, Line: 1, Column: 0, Text: "\"", Value: ""},
			Token{Type: TokenTypeEOF, Line: 1, Column: 1, Text: ``},
		},
	},
	{
		src: "[]{},.:;+-=*&()<>!#",
		results: []Token{
			Token{Type: TokenTypeLBRACK, Column: 0, Text: `[`},
			Token{Type: TokenTypeRBRACK, Column: 1, Text: `]`},
			Token{Type: TokenTypeLBRACE, Column: 2, Text: `{`},
			Token{Type: TokenTypeRBRACE, Column: 3, Text: `}`},
			Token{Type: TokenTypeCOMMA, Column: 4, Text: `,`},
			Token{Type: TokenTypePERIOD, Column: 5, Text: `.`},
			Token{Type: TokenTypeCOLON, Column: 6, Text: `:`},
			Token{Type: TokenTypeSEMICOLON, Column: 7, Text: `;`},
			Token{Type: TokenTypeADD, Column: 8, Text: `+`},
			Token{Type: TokenTypeSUB, Column: 9, Text: `-`},
			Token{Type: TokenTypeASSIGN, Column: 10, Text: `=`},
			Token{Type: TokenTypeASTERIX, Column: 11, Text: `*`},
			Token{Type: TokenTypeAMPERSAND, Column: 12, Text: `&`},
			Token{Type: TokenTypeLPAREN, Column: 13, Text: `(`},
			Token{Type: TokenTypeRPAREN, Column: 14, Text: `)`},
			Token{Type: TokenTypeLCHEV, Column: 15, Text: `<`},
			Token{Type: TokenTypeRCHEV, Column: 16, Text: `>`},
			Token{Type: TokenTypeEXCL, Column: 17, Text: `!`},
			Token{Type: TokenTypeHASHBANG, Column: 18, Text: `#`},
			Token{Type: TokenTypeEOF, Column: 19, Text: ``},
		},
	},
	{
		src: "false true",
		results: []Token{
			Token{Type: TokenTypeBoolean, Column: 0, Text: `false`, Value: false},
			Token{Type: TokenTypeWhitespace, Column: 5, Text: ` `},
			Token{Type: TokenTypeBoolean, Column: 6, Text: `true`, Value: true},
			Token{Type: TokenTypeEOF, Column: 10, Text: ``},
		},
	},
	{
		src: "12348 1234.5 1234.5.5", // Sounds like a hack but I'll leave it to the parser to decide what to do
		results: []Token{
			Token{Type: TokenTypeNumber, Column: 0, Text: `12348`, Value: int64(12348)},
			Token{Type: TokenTypeWhitespace, Column: 5, Text: ` `},
			Token{Type: TokenTypeFloat, Column: 6, Text: `1234.5`, Value: 1234.5},
			Token{Type: TokenTypeWhitespace, Column: 12, Text: ` `},
			Token{Type: TokenTypeFloat, Column: 13, Text: `1234.5`, Value: 1234.5},
			Token{Type: TokenTypeFloat, Column: 19, Text: `.5`, Value: 0.5},
			Token{Type: TokenTypeEOF, Column: 21, Text: ``},
		},
	},
	{
		src: "// This is a comment\n/*\nfoo\n*/// eof comment",
		results: []Token{
			Token{Type: TokenTypeComment, Column: 0, Text: `// This is a comment`},
			Token{Type: TokenTypeWhitespace, Column: 20, Text: "\n"},
			Token{Type: TokenTypeComment, Column: 0, Line: 1, Text: "/*\nfoo\n*/"},
			Token{Type: TokenTypeComment, Column: 2, Line: 3, Text: `// eof comment`},
			Token{Type: TokenTypeEOF, Column: 16, Line: 3},
		},
	},
	{
		src: "/* eof ending block comments wont work",
		results: []Token{
			Token{Type: TokenTypeUnknown, Column: 0, Text: `/* eof ending block comments wont work`},
			Token{Type: TokenTypeEOF, Column: 38, Line: 0},
		},
	},
}

func TestScannerTable(t *testing.T) {
	for _, test := range tests {
		s := NewScanner(strings.NewReader(test.src))
		tokens := []Token{}
		for token := range s.ScanChannel() {
			tokens = append(tokens, token)
			if token.Type == TokenTypeEOF {
				break
			}
		}

		once := sync.Once{}

		for i, token := range tokens {
			var expectedToken Token
			if len(test.results) > i {
				expectedToken = test.results[i]
			} else {
				t.Error("Too many token returned")
			}
			if !reflect.DeepEqual(expectedToken, token) {
				once.Do(func() {
					t.Error("source", test.src)
				})
				t.Errorf("Token (%d): %#v didn't match expected %#v", i, token, expectedToken)
			}
		}

	}
}

func TestTokenizesC(t *testing.T) {
	s := NewScanner(strings.NewReader(`
    #include <stdio.h>
    int main()
    {
      // printf() displays the string inside quotation
      printf("Hello, World!");
      return 0;
    }
  `))

	for token := range s.ScanChannel() {
		if token.Type == TokenTypeUnknown {
			t.Errorf("Encountered an unknown token %s", token)
		}
	}
}

func TestTokenizesGO(t *testing.T) {
	s := NewScanner(strings.NewReader(`
    package main
    import "fmt"
    func main() {
      fmt.Println("hello world")
    }
  `))

	for token := range s.ScanChannel() {
		if token.Type == TokenTypeUnknown {
			t.Errorf("Encountered an unknown token %s", token)
		}
	}
}

func ExampleNewScanner() {
	s := NewScanner(strings.NewReader(`
    fn test(foo: Int, bar: Test) Foo { // 1
      var a = foo; // 2
      var b = 123; // 3
      var c = true; // 4
      // This is a comment 5
      var d = 123.5; // 6
      /* One line block comment */ // 7
      /* 8
       * Multiline comment 9
       */ // 10
      return "foo"; // 11
    } // 12
  `))
	tokens := []string{}
	for token := range s.ScanChannel() {
		if token.Type != TokenTypeWhitespace {
			tokens = append(tokens, strings.Replace(token.String(), "\n", "[NL]", -1))
		}
	}

	fmt.Print(strings.Join(tokens, "\n"))
	// Output:
	// 1:4 IDENT(fn)
	// 1:7 IDENT(test)
	// 1:11 LPAREN(()
	// 1:12 IDENT(foo)
	// 1:15 COLON(:)
	// 1:17 IDENT(Int)
	// 1:20 COMMA(,)
	// 1:22 IDENT(bar)
	// 1:25 COLON(:)
	// 1:27 IDENT(Test)
	// 1:31 RPAREN())
	// 1:33 IDENT(Foo)
	// 1:37 LBRACE({)
	// 1:39 COMMENT(// 1)
	// 2:6 IDENT(var)
	// 2:10 IDENT(a)
	// 2:12 ASSIGN(=)
	// 2:14 IDENT(foo)
	// 2:17 SEMICOLON(;)
	// 2:19 COMMENT(// 2)
	// 3:6 IDENT(var)
	// 3:10 IDENT(b)
	// 3:12 ASSIGN(=)
	// 3:14 NUMBER(123)
	// 3:17 SEMICOLON(;)
	// 3:19 COMMENT(// 3)
	// 4:6 IDENT(var)
	// 4:10 IDENT(c)
	// 4:12 ASSIGN(=)
	// 4:14 BOOL(true)
	// 4:18 SEMICOLON(;)
	// 4:20 COMMENT(// 4)
	// 5:6 COMMENT(// This is a comment 5)
	// 6:6 IDENT(var)
	// 6:10 IDENT(d)
	// 6:12 ASSIGN(=)
	// 6:14 FLOAT(123.5)
	// 6:19 SEMICOLON(;)
	// 6:21 COMMENT(// 6)
	// 7:6 COMMENT(/* One line block comment */)
	// 7:35 COMMENT(// 7)
	// 8:6 COMMENT(/* 8[NL]       * Multiline comment 9[NL]       */)
	// 10:10 COMMENT(// 10)
	// 11:6 IDENT(return)
	// 11:13 STRING(foo)
	// 11:18 SEMICOLON(;)
	// 11:20 COMMENT(// 11)
	// 12:4 RBRACE(})
	// 12:6 COMMENT(// 12)
	// 13:2 EOF()
}
<commit_msg>Benchmarks and tests<commit_after>package scanner

import (
	"fmt"
	"reflect"
	"strings"
	"sync"
	"testing"
)

var tests = []struct {
	src     string
	results []Token
}{
	{
		src: "export fn\nfn1",
		results: []Token{
			Token{Type: TokenTypeIdent, Text: "export"},
			Token{Type: TokenTypeWhitespace, Column: 6, Text: " "},
			Token{Type: TokenTypeIdent, Column: 7, Text: "fn"},
			Token{Type: TokenTypeWhitespace, Column: 9, Text: "\n"},
			Token{Type: TokenTypeIdent, Line: 1, Text: "fn1"},
			Token{Type: TokenTypeEOF, Line: 1, Column: 3, Text: ""},
		},
	},
	{
		src: "\"foobar\"\n\"foo",
		results: []Token{
			Token{Type: TokenTypeString, Text: `"foobar"`, Value: "foobar"},
			Token{Type: TokenTypeWhitespace, Column: 8, Text: "\n"},
			Token{Type: TokenTypeUnknown, Line: 1, Column: 0, Text: `"foo`, Value: ""},
			Token{Type: TokenTypeEOF, Line: 1, Column: 4, Text: ``},
		},
	},
	{
		src: `"foo\"bar""foo\"bar"`,
		results: []Token{
			Token{Type: TokenTypeString, Text: `"foo\"bar"`, Value: `foo"bar`},
			Token{Type: TokenTypeString, Column: 10, Text: `"foo\"bar"`, Value: `foo"bar`},
			Token{Type: TokenTypeEOF, Line: 0, Column: 20, Text: ``},
		},
	},
	{
		src: "\"foo\\",
		results: []Token{
			Token{Type: TokenTypeUnknown, Text: `"foo\`, Value: ""},
			Token{Type: TokenTypeEOF, Line: 0, Column: 5, Text: ``},
		},
	},
	{
		src: `"foo\nbar"`,
		results: []Token{
			Token{Type: TokenTypeString, Text: `"foo\nbar"`, Value: "foo\\nbar"},
			Token{Type: TokenTypeEOF, Line: 0, Column: 10, Text: ``},
		},
	},
	{
		src: `"foo\\bar"`,
		results: []Token{
			Token{Type: TokenTypeString, Text: `"foo\\bar"`, Value: `foo\bar`},
			Token{Type: TokenTypeEOF, Line: 0, Column: 10, Text: ``},
		},
	},
	{
		src: "\"\\123\\x53\\u2318\"",
		results: []Token{
			Token{Type: TokenTypeString, Text: "\"\\123\\x53\\u2318\"", Value: "SS⌘"},
			Token{Type: TokenTypeEOF, Line: 0, Column: 16, Text: ``},
		},
	},
	/*
		{
			src: `"\xe6\x97\xa5\xe6\x9c\xac\xe8\xaa\x9e"`,
			results: []Token{
				Token{Type: TokenTypeString, Text: `"\xe6\x97\xa5\xe6\x9c\xac\xe8\xaa\x9e"`, Value: "日本語"},
				Token{Type: TokenTypeEOF, Line: 0, Column: 38, Text: ``},
			},
		},
	*/
	{
		src: "`\\n\\000`",
		results: []Token{
			Token{Type: TokenTypeString, Text: "`\\n\\000`", Value: "\\n\\000"},
			Token{Type: TokenTypeEOF, Line: 0, Column: 8, Text: ``},
		},
	},
	{
		src: "\"\n\"",
		results: []Token{
			Token{Type: TokenTypeUnknown, Text: "\"", Value: ""},
			Token{Type: TokenTypeWhitespace, Line: 0, Column: 1, Text: "\n", Value: nil},
			Token{Type: TokenTypeUnknown, Line: 1, Column: 0, Text: "\"", Value: ""},
			Token{Type: TokenTypeEOF, Line: 1, Column: 1, Text: ``},
		},
	},
	{
		src: "[]{},.:;+-=*&()<>!#",
		results: []Token{
			Token{Type: TokenTypeLBRACK, Column: 0, Text: `[`},
			Token{Type: TokenTypeRBRACK, Column: 1, Text: `]`},
			Token{Type: TokenTypeLBRACE, Column: 2, Text: `{`},
			Token{Type: TokenTypeRBRACE, Column: 3, Text: `}`},
			Token{Type: TokenTypeCOMMA, Column: 4, Text: `,`},
			Token{Type: TokenTypePERIOD, Column: 5, Text: `.`},
			Token{Type: TokenTypeCOLON, Column: 6, Text: `:`},
			Token{Type: TokenTypeSEMICOLON, Column: 7, Text: `;`},
			Token{Type: TokenTypeADD, Column: 8, Text: `+`},
			Token{Type: TokenTypeSUB, Column: 9, Text: `-`},
			Token{Type: TokenTypeASSIGN, Column: 10, Text: `=`},
			Token{Type: TokenTypeASTERIX, Column: 11, Text: `*`},
			Token{Type: TokenTypeAMPERSAND, Column: 12, Text: `&`},
			Token{Type: TokenTypeLPAREN, Column: 13, Text: `(`},
			Token{Type: TokenTypeRPAREN, Column: 14, Text: `)`},
			Token{Type: TokenTypeLCHEV, Column: 15, Text: `<`},
			Token{Type: TokenTypeRCHEV, Column: 16, Text: `>`},
			Token{Type: TokenTypeEXCL, Column: 17, Text: `!`},
			Token{Type: TokenTypeHASHBANG, Column: 18, Text: `#`},
			Token{Type: TokenTypeEOF, Column: 19, Text: ``},
		},
	},
	{
		src: "false true",
		results: []Token{
			Token{Type: TokenTypeBoolean, Column: 0, Text: `false`, Value: false},
			Token{Type: TokenTypeWhitespace, Column: 5, Text: ` `},
			Token{Type: TokenTypeBoolean, Column: 6, Text: `true`, Value: true},
			Token{Type: TokenTypeEOF, Column: 10, Text: ``},
		},
	},
	{
		src: "12348 1234.5 1234.5.5", // Sounds like a hack but I'll leave it to the parser to decide what to do
		results: []Token{
			Token{Type: TokenTypeNumber, Column: 0, Text: `12348`, Value: int64(12348)},
			Token{Type: TokenTypeWhitespace, Column: 5, Text: ` `},
			Token{Type: TokenTypeFloat, Column: 6, Text: `1234.5`, Value: 1234.5},
			Token{Type: TokenTypeWhitespace, Column: 12, Text: ` `},
			Token{Type: TokenTypeFloat, Column: 13, Text: `1234.5`, Value: 1234.5},
			Token{Type: TokenTypeFloat, Column: 19, Text: `.5`, Value: 0.5},
			Token{Type: TokenTypeEOF, Column: 21, Text: ``},
		},
	},
	{
		src: "// This is a comment\n/*\nfoo\n*/// eof comment",
		results: []Token{
			Token{Type: TokenTypeComment, Column: 0, Text: `// This is a comment`},
			Token{Type: TokenTypeWhitespace, Column: 20, Text: "\n"},
			Token{Type: TokenTypeComment, Column: 0, Line: 1, Text: "/*\nfoo\n*/"},
			Token{Type: TokenTypeComment, Column: 2, Line: 3, Text: `// eof comment`},
			Token{Type: TokenTypeEOF, Column: 16, Line: 3},
		},
	},
	{
		src: "/* eof ending block comments wont work",
		results: []Token{
			Token{Type: TokenTypeUnknown, Column: 0, Text: `/* eof ending block comments wont work`},
			Token{Type: TokenTypeEOF, Column: 38, Line: 0},
		},
	},
}

func TestScannerTable(t *testing.T) {
	for _, test := range tests {
		s := NewScanner(strings.NewReader(test.src))
		tokens := []Token{}
		for token := range s.ScanChannel() {
			tokens = append(tokens, token)
			if token.Type == TokenTypeEOF {
				break
			}
		}

		once := sync.Once{}

		for i, token := range tokens {
			var expectedToken Token
			if len(test.results) > i {
				expectedToken = test.results[i]
			} else {
				t.Error("Too many token returned")
			}
			if !reflect.DeepEqual(expectedToken, token) {
				once.Do(func() {
					t.Error("source", test.src)
				})
				t.Errorf("Token (%d): %#v didn't match expected %#v", i, token, expectedToken)
			}
		}

	}
}

func BenchmarkScannerTable(b *testing.B) {
	for i := 0; i  i {
					expectedToken = test.results[i]
				} else {
					b.Error("Too many token returned")
				}
				if !reflect.DeepEqual(expectedToken, token) {
					once.Do(func() {
						b.Error("source", test.src)
					})
					b.Errorf("Token (%d): %#v didn't match expected %#v", i, token, expectedToken)
				}
			}

		}
	}
}

func TestTokenizesC(t *testing.T) {
	s := NewScanner(strings.NewReader(`
    #include <stdio.h>
    int main()
    {
      // printf() displays the string inside quotation
      printf("Hello, World!");
      return 0;
    }
  `))

	for token := range s.ScanChannel() {
		if token.Type == TokenTypeUnknown {
			t.Errorf("Encountered an unknown token %s", token)
		}
	}
}

func BenchmarkTokenizesC(b *testing.B) {
	for i := 0; i < b.N; i++ {
		s := NewScanner(strings.NewReader(`
    #include <stdio.h>
    int main()
    {
      // printf() displays the string inside quotation
      printf("Hello, World!");
      return 0;
    }
  `))

		for token := range s.ScanChannel() {
			if token.Type == TokenTypeUnknown {
				b.Errorf("Encountered an unknown token %s", token)
			}
		}
	}
}

func TestTokenizesGO(t *testing.T) {
	s := NewScanner(strings.NewReader(`
    package main
    import "fmt"
    func main() {
      fmt.Println("hello world")
    }
  `))

	for token := range s.ScanChannel() {
		if token.Type == TokenTypeUnknown {
			t.Errorf("Encountered an unknown token %s", token)
		}
	}
}

func TestTokenizesJSX(t *testing.T) {
	s := NewScanner(strings.NewReader(`

TODO

        <TodoList items={this.state.items} />
        <form onSubmit={this.handleSubmit}>
          <input onChange={this.handleChange} value={this.state.text} />
          <button>{'Add #' + (this.state.items.length + 1)}</button>
        </form>

  `))

	for token := range s.ScanChannel() {
		println(token.String())
		if token.Type == TokenTypeUnknown {
			t.Errorf("Encountered an unknown token %s", token)
		}
	}
}

func ExampleNewScanner() {
	s := NewScanner(strings.NewReader(`
    fn test(foo: Int, bar: Test) Foo { // 1
      var a = foo; // 2
      var b = 123; // 3
      var c = true; // 4
      // This is a comment 5
      var d = 123.5; // 6
      /* One line block comment */ // 7
      /* 8
       * Multiline comment 9
       */ // 10
      return "foo"; // 11
    } // 12
  `))
	tokens := []string{}
	for token := range s.ScanChannel() {
		if token.Type != TokenTypeWhitespace {
			tokens = append(tokens, strings.Replace(token.String(), "\n", "[NL]", -1))
		}
	}

	fmt.Print(strings.Join(tokens, "\n"))
	// Output:
	// 1:4 IDENT(fn)
	// 1:7 IDENT(test)
	// 1:11 LPAREN(()
	// 1:12 IDENT(foo)
	// 1:15 COLON(:)
	// 1:17 IDENT(Int)
	// 1:20 COMMA(,)
	// 1:22 IDENT(bar)
	// 1:25 COLON(:)
	// 1:27 IDENT(Test)
	// 1:31 RPAREN())
	// 1:33 IDENT(Foo)
	// 1:37 LBRACE({)
	// 1:39 COMMENT(// 1)
	// 2:6 IDENT(var)
	// 2:10 IDENT(a)
	// 2:12 ASSIGN(=)
	// 2:14 IDENT(foo)
	// 2:17 SEMICOLON(;)
	// 2:19 COMMENT(// 2)
	// 3:6 IDENT(var)
	// 3:10 IDENT(b)
	// 3:12 ASSIGN(=)
	// 3:14 NUMBER(123)
	// 3:17 SEMICOLON(;)
	// 3:19 COMMENT(// 3)
	// 4:6 IDENT(var)
	// 4:10 IDENT(c)
	// 4:12 ASSIGN(=)
	// 4:14 BOOL(true)
	// 4:18 SEMICOLON(;)
	// 4:20 COMMENT(// 4)
	// 5:6 COMMENT(// This is a comment 5)
	// 6:6 IDENT(var)
	// 6:10 IDENT(d)
	// 6:12 ASSIGN(=)
	// 6:14 FLOAT(123.5)
	// 6:19 SEMICOLON(;)
	// 6:21 COMMENT(// 6)
	// 7:6 COMMENT(/* One line block comment */)
	// 7:35 COMMENT(// 7)
	// 8:6 COMMENT(/* 8[NL]       * Multiline comment 9[NL]       */)
	// 10:10 COMMENT(// 10)
	// 11:6 IDENT(return)
	// 11:13 STRING(foo)
	// 11:18 SEMICOLON(;)
	// 11:20 COMMENT(// 11)
	// 12:4 RBRACE(})
	// 12:6 COMMENT(// 12)
	// 13:2 EOF()
}
<|endoftext|>


<commit_before>package session

import (
	"os"
	"strconv"

	"github.com/aws/aws-sdk-go/aws/credentials"
)

// envConfig is a collection of environment values the SDK will read
// setup config from. All environment values are optional. But some values
// such as credentials require multiple values to be complete or the values
// will be ignored.
type envConfig struct {
	// Environment configuration values. If set both Access Key ID and Secret Access
	// Key must be provided. Session Token and optionally also be provided, but is
	// not required.
	//
	//	# Access Key ID
	//	AWS_ACCESS_KEY_ID=AKID
	//	AWS_ACCESS_KEY=AKID # only read if AWS_ACCESS_KEY_ID is not set.
	//
	//	# Secret Access Key
	//	AWS_SECRET_ACCESS_KEY=SECRET
	//	AWS_SECRET_KEY=SECRET=SECRET # only read if AWS_SECRET_ACCESS_KEY is not set.
	//
	//	# Session Token
	//	AWS_SESSION_TOKEN=TOKEN
	Creds credentials.Value

	// Region value will instruct the SDK where to make service API requests to. If is
	// not provided in the environment the region must be provided before a service
	// client request is made.
	//
	//	AWS_REGION=us-east-1
	//
	//	# AWS_DEFAULT_REGION is only read if AWS_SDK_LOAD_CONFIG is also set,
	//	# and AWS_REGION is not also set.
	//	AWS_DEFAULT_REGION=us-east-1
	Region string

	// Profile name the SDK should load use when loading shared configuration from the
	// shared configuration files. If not provided "default" will be used as the
	// profile name.
	//
	//	AWS_PROFILE=my_profile
	//
	//	# AWS_DEFAULT_PROFILE is only read if AWS_SDK_LOAD_CONFIG is also set,
	//	# and AWS_PROFILE is not also set.
	//	AWS_DEFAULT_PROFILE=my_profile
	Profile string

	// SDK load config instructs the SDK to load the shared config in addition to
	// shared credentials. This also expands the configuration loaded from the shared
	// credentials to have parity with the shared config file. This also enables
	// Region and Profile support for the AWS_DEFAULT_REGION and AWS_DEFAULT_PROFILE
	// env values as well.
	//
	//	AWS_SDK_LOAD_CONFIG=1
	EnableSharedConfig bool

	// Shared credentials file path can be set to instruct the SDK to use an alternate
	// file for the shared credentials. If not set the file will be loaded from
	// $HOME/.aws/credentials on Linux/Unix based systems, and
	// %USERPROFILE%\.aws\credentials on Windows.
	//
	//	AWS_SHARED_CREDENTIALS_FILE=$HOME/my_shared_credentials
	SharedCredentialsFile string

	// Shared config file path can be set to instruct the SDK to use an alternate
	// file for the shared config. If not set the file will be loaded from
	// $HOME/.aws/config on Linux/Unix based systems, and
	// %USERPROFILE%\.aws\config on Windows.
	//
	//	AWS_CONFIG_FILE=$HOME/my_shared_config
	SharedConfigFile string

	// Sets the path to a custom Credentials Authroity (CA) Bundle PEM file
	// that the SDK will use instead of the system's root CA bundle.
	// Only use this if you want to configure the SDK to use a custom set
	// of CAs.
	//
	// Enabling this option will attempt to merge the Transport
	// into the SDK's HTTP client. If the client's Transport is
	// not a http.Transport an error will be returned. If the
	// Transport's TLS config is set this option will cause the
	// SDK to overwrite the Transport's TLS config's  RootCAs value.
	//
	// Setting a custom HTTPClient in the aws.Config options will override this setting.
	// To use this option and custom HTTP client, the HTTP client needs to be provided
	// when creating the session. Not the service client.
	//
	//  AWS_CA_BUNDLE=$HOME/my_custom_ca_bundle
	CustomCABundle string
}

var (
	credAccessEnvKey = []string{
		"AWS_ACCESS_KEY_ID",
		"AWS_ACCESS_KEY",
	}
	credSecretEnvKey = []string{
		"AWS_SECRET_ACCESS_KEY",
		"AWS_SECRET_KEY",
	}
	credSessionEnvKey = []string{
		"AWS_SESSION_TOKEN",
	}

	regionEnvKeys = []string{
		"AWS_REGION",
		"AWS_DEFAULT_REGION", // Only read if AWS_SDK_LOAD_CONFIG is also set
	}
	profileEnvKeys = []string{
		"AWS_PROFILE",
		"AWS_DEFAULT_PROFILE", // Only read if AWS_SDK_LOAD_CONFIG is also set
	}
	sharedCredsFileEnvKey = []string{
		"AWS_SHARED_CREDENTIALS_FILE",
	}
	sharedConfigFileEnvKey = []string{
		"AWS_CONFIG_FILE",
	}
)

// loadEnvConfig retrieves the SDK's environment configuration.
// See `envConfig` for the values that will be retrieved.
//
// If the environment variable `AWS_SDK_LOAD_CONFIG` is set to a truthy value
// the shared SDK config will be loaded in addition to the SDK's specific
// configuration values.
func loadEnvConfig() envConfig {
	enableSharedConfig, _ := strconv.ParseBool(os.Getenv("AWS_SDK_LOAD_CONFIG"))
	return envConfigLoad(enableSharedConfig)
}

// loadEnvSharedConfig retrieves the SDK's environment configuration, and the
// SDK shared config. See `envConfig` for the values that will be retrieved.
//
// Loads the shared configuration in addition to the SDK's specific configuration.
// This will load the same values as `loadEnvConfig` if the `AWS_SDK_LOAD_CONFIG`
// environment variable is set.
func loadSharedEnvConfig() envConfig {
	return envConfigLoad(true)
}

func envConfigLoad(enableSharedConfig bool) envConfig {
	cfg := envConfig{}

	cfg.EnableSharedConfig = enableSharedConfig

	setFromEnvVal(&cfg.Creds.AccessKeyID, credAccessEnvKey)
	setFromEnvVal(&cfg.Creds.SecretAccessKey, credSecretEnvKey)
	setFromEnvVal(&cfg.Creds.SessionToken, credSessionEnvKey)

	// Require logical grouping of credentials
	if len(cfg.Creds.AccessKeyID) == 0 || len(cfg.Creds.SecretAccessKey) == 0 {
		cfg.Creds = credentials.Value{}
	} else {
		cfg.Creds.ProviderName = "EnvConfigCredentials"
	}

	regionKeys := regionEnvKeys
	profileKeys := profileEnvKeys
	if !cfg.EnableSharedConfig {
		regionKeys = regionKeys[:1]
		profileKeys = profileKeys[:1]
	}

	setFromEnvVal(&cfg.Region, regionKeys)
	setFromEnvVal(&cfg.Profile, profileKeys)

	setFromEnvVal(&cfg.SharedCredentialsFile, sharedCredsFileEnvKey)
	setFromEnvVal(&cfg.SharedConfigFile, sharedConfigFileEnvKey)

	cfg.CustomCABundle = os.Getenv("AWS_CA_BUNDLE")

	return cfg
}

func setFromEnvVal(dst *string, keys []string) {
	for _, k := range keys {
		if v := os.Getenv(k); len(v) > 0 {
			*dst = v
			break
		}
	}
}
<commit_msg>aws/credentials: Add EnvProviderName constant. (#1531)<commit_after>package session

import (
	"os"
	"strconv"

	"github.com/aws/aws-sdk-go/aws/credentials"
)

// EnvProviderName provides a name of the provider when config is loaded from environment.
const EnvProviderName = "EnvConfigCredentials"

// envConfig is a collection of environment values the SDK will read
// setup config from. All environment values are optional. But some values
// such as credentials require multiple values to be complete or the values
// will be ignored.
type envConfig struct {
	// Environment configuration values. If set both Access Key ID and Secret Access
	// Key must be provided. Session Token and optionally also be provided, but is
	// not required.
	//
	//	# Access Key ID
	//	AWS_ACCESS_KEY_ID=AKID
	//	AWS_ACCESS_KEY=AKID # only read if AWS_ACCESS_KEY_ID is not set.
	//
	//	# Secret Access Key
	//	AWS_SECRET_ACCESS_KEY=SECRET
	//	AWS_SECRET_KEY=SECRET=SECRET # only read if AWS_SECRET_ACCESS_KEY is not set.
	//
	//	# Session Token
	//	AWS_SESSION_TOKEN=TOKEN
	Creds credentials.Value

	// Region value will instruct the SDK where to make service API requests to. If is
	// not provided in the environment the region must be provided before a service
	// client request is made.
	//
	//	AWS_REGION=us-east-1
	//
	//	# AWS_DEFAULT_REGION is only read if AWS_SDK_LOAD_CONFIG is also set,
	//	# and AWS_REGION is not also set.
	//	AWS_DEFAULT_REGION=us-east-1
	Region string

	// Profile name the SDK should load use when loading shared configuration from the
	// shared configuration files. If not provided "default" will be used as the
	// profile name.
	//
	//	AWS_PROFILE=my_profile
	//
	//	# AWS_DEFAULT_PROFILE is only read if AWS_SDK_LOAD_CONFIG is also set,
	//	# and AWS_PROFILE is not also set.
	//	AWS_DEFAULT_PROFILE=my_profile
	Profile string

	// SDK load config instructs the SDK to load the shared config in addition to
	// shared credentials. This also expands the configuration loaded from the shared
	// credentials to have parity with the shared config file. This also enables
	// Region and Profile support for the AWS_DEFAULT_REGION and AWS_DEFAULT_PROFILE
	// env values as well.
	//
	//	AWS_SDK_LOAD_CONFIG=1
	EnableSharedConfig bool

	// Shared credentials file path can be set to instruct the SDK to use an alternate
	// file for the shared credentials. If not set the file will be loaded from
	// $HOME/.aws/credentials on Linux/Unix based systems, and
	// %USERPROFILE%\.aws\credentials on Windows.
	//
	//	AWS_SHARED_CREDENTIALS_FILE=$HOME/my_shared_credentials
	SharedCredentialsFile string

	// Shared config file path can be set to instruct the SDK to use an alternate
	// file for the shared config. If not set the file will be loaded from
	// $HOME/.aws/config on Linux/Unix based systems, and
	// %USERPROFILE%\.aws\config on Windows.
	//
	//	AWS_CONFIG_FILE=$HOME/my_shared_config
	SharedConfigFile string

	// Sets the path to a custom Credentials Authroity (CA) Bundle PEM file
	// that the SDK will use instead of the system's root CA bundle.
	// Only use this if you want to configure the SDK to use a custom set
	// of CAs.
	//
	// Enabling this option will attempt to merge the Transport
	// into the SDK's HTTP client. If the client's Transport is
	// not a http.Transport an error will be returned. If the
	// Transport's TLS config is set this option will cause the
	// SDK to overwrite the Transport's TLS config's  RootCAs value.
	//
	// Setting a custom HTTPClient in the aws.Config options will override this setting.
	// To use this option and custom HTTP client, the HTTP client needs to be provided
	// when creating the session. Not the service client.
	//
	//  AWS_CA_BUNDLE=$HOME/my_custom_ca_bundle
	CustomCABundle string
}

var (
	credAccessEnvKey = []string{
		"AWS_ACCESS_KEY_ID",
		"AWS_ACCESS_KEY",
	}
	credSecretEnvKey = []string{
		"AWS_SECRET_ACCESS_KEY",
		"AWS_SECRET_KEY",
	}
	credSessionEnvKey = []string{
		"AWS_SESSION_TOKEN",
	}

	regionEnvKeys = []string{
		"AWS_REGION",
		"AWS_DEFAULT_REGION", // Only read if AWS_SDK_LOAD_CONFIG is also set
	}
	profileEnvKeys = []string{
		"AWS_PROFILE",
		"AWS_DEFAULT_PROFILE", // Only read if AWS_SDK_LOAD_CONFIG is also set
	}
	sharedCredsFileEnvKey = []string{
		"AWS_SHARED_CREDENTIALS_FILE",
	}
	sharedConfigFileEnvKey = []string{
		"AWS_CONFIG_FILE",
	}
)

// loadEnvConfig retrieves the SDK's environment configuration.
// See `envConfig` for the values that will be retrieved.
//
// If the environment variable `AWS_SDK_LOAD_CONFIG` is set to a truthy value
// the shared SDK config will be loaded in addition to the SDK's specific
// configuration values.
func loadEnvConfig() envConfig {
	enableSharedConfig, _ := strconv.ParseBool(os.Getenv("AWS_SDK_LOAD_CONFIG"))
	return envConfigLoad(enableSharedConfig)
}

// loadEnvSharedConfig retrieves the SDK's environment configuration, and the
// SDK shared config. See `envConfig` for the values that will be retrieved.
//
// Loads the shared configuration in addition to the SDK's specific configuration.
// This will load the same values as `loadEnvConfig` if the `AWS_SDK_LOAD_CONFIG`
// environment variable is set.
func loadSharedEnvConfig() envConfig {
	return envConfigLoad(true)
}

func envConfigLoad(enableSharedConfig bool) envConfig {
	cfg := envConfig{}

	cfg.EnableSharedConfig = enableSharedConfig

	setFromEnvVal(&cfg.Creds.AccessKeyID, credAccessEnvKey)
	setFromEnvVal(&cfg.Creds.SecretAccessKey, credSecretEnvKey)
	setFromEnvVal(&cfg.Creds.SessionToken, credSessionEnvKey)

	// Require logical grouping of credentials
	if len(cfg.Creds.AccessKeyID) == 0 || len(cfg.Creds.SecretAccessKey) == 0 {
		cfg.Creds = credentials.Value{}
	} else {
		cfg.Creds.ProviderName = EnvProviderName
	}

	regionKeys := regionEnvKeys
	profileKeys := profileEnvKeys
	if !cfg.EnableSharedConfig {
		regionKeys = regionKeys[:1]
		profileKeys = profileKeys[:1]
	}

	setFromEnvVal(&cfg.Region, regionKeys)
	setFromEnvVal(&cfg.Profile, profileKeys)

	setFromEnvVal(&cfg.SharedCredentialsFile, sharedCredsFileEnvKey)
	setFromEnvVal(&cfg.SharedConfigFile, sharedConfigFileEnvKey)

	cfg.CustomCABundle = os.Getenv("AWS_CA_BUNDLE")

	return cfg
}

func setFromEnvVal(dst *string, keys []string) {
	for _, k := range keys {
		if v := os.Getenv(k); len(v) > 0 {
			*dst = v
			break
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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 gce

import "k8s.io/api/core/v1"

type LoadBalancerType string

const (
	// ServiceAnnotationLoadBalancerType is annotated on a service with type LoadBalancer
	// dictates what specific kind of GCP LB should be assembled.
	// Currently, only "internal" is supported.
	ServiceAnnotationLoadBalancerType = "cloud.google.com/load-balancer-type"

	LBTypeInternal LoadBalancerType = "internal"

	// ServiceAnnotationInternalBackendShare is annotated on a service with "true" when users
	// want to share GCP Backend Services for a set of internal load balancers.
	// ALPHA feature - this may be removed in a future release.
	ServiceAnnotationILBBackendShare = "cloud.google.com/load-balancer-backend-share"
)

// GetLoadBalancerAnnotationType returns the type of GCP load balancer which should be assembled.
func GetLoadBalancerAnnotationType(service *v1.Service) (LoadBalancerType, bool) {
	v := LoadBalancerType("")
	if service.Spec.Type != v1.ServiceTypeLoadBalancer {
		return v, false
	}

	l, ok := service.Annotations[ServiceAnnotationLoadBalancerType]
	v = LoadBalancerType(l)
	if !ok {
		return v, false
	}

	switch v {
	case LBTypeInternal:
		return v, true
	default:
		return v, false
	}
}

// GetLoadBalancerAnnotationBackendShare returns whether this service's backend service should be
// shared with other load balancers. Health checks and the healthcheck firewall will be shared regardless.
func GetLoadBalancerAnnotationBackendShare(service *v1.Service) bool {
	l, exists := service.Annotations[ServiceAnnotationILBBackendShare]
	if exists && l == "true" {
		return true
	}

	return false
}
<commit_msg>GCE: Specify alpha in annotation key, deprecate lower case of LB type<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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 gce

import (
	"github.com/golang/glog"
	"k8s.io/api/core/v1"
)

type LoadBalancerType string

const (
	// ServiceAnnotationLoadBalancerType is annotated on a service with type LoadBalancer
	// dictates what specific kind of GCP LB should be assembled.
	// Currently, only "internal" is supported.
	ServiceAnnotationLoadBalancerType = "cloud.google.com/load-balancer-type"

	LBTypeInternal LoadBalancerType = "Internal"
	// Deprecating the lowercase spelling of Internal.
	deprecatedTypeInternalLowerCase LoadBalancerType = "internal"

	// ServiceAnnotationInternalBackendShare is annotated on a service with "true" when users
	// want to share GCP Backend Services for a set of internal load balancers.
	// ALPHA feature - this may be removed in a future release.
	ServiceAnnotationILBBackendShare = "alpha.cloud.google.com/load-balancer-backend-share"
	// This annotation did not correctly specify "alpha", so both annotations will be checked.
	deprecatedServiceAnnotationILBBackendShare = "cloud.google.com/load-balancer-backend-share"
)

// GetLoadBalancerAnnotationType returns the type of GCP load balancer which should be assembled.
func GetLoadBalancerAnnotationType(service *v1.Service) (LoadBalancerType, bool) {
	v := LoadBalancerType("")
	if service.Spec.Type != v1.ServiceTypeLoadBalancer {
		return v, false
	}

	l, ok := service.Annotations[ServiceAnnotationLoadBalancerType]
	v = LoadBalancerType(l)
	if !ok {
		return v, false
	}

	switch v {
	case LBTypeInternal, deprecatedTypeInternalLowerCase:
		return LBTypeInternal, true
	default:
		return v, false
	}
}

// GetLoadBalancerAnnotationBackendShare returns whether this service's backend service should be
// shared with other load balancers. Health checks and the healthcheck firewall will be shared regardless.
func GetLoadBalancerAnnotationBackendShare(service *v1.Service) bool {
	if l, exists := service.Annotations[ServiceAnnotationILBBackendShare]; exists && l == "true" {
		return true
	}

	// Check for deprecated annotation key
	if l, exists := service.Annotations[deprecatedServiceAnnotationILBBackendShare]; exists && l == "true" {
		glog.Warningf("Annotation %q is deprecated and replaced with an alpha-specific key: %q", deprecatedServiceAnnotationILBBackendShare, ServiceAnnotationILBBackendShare)
		return true
	}

	return false
}
<|endoftext|>


<commit_before>// Copyright 2022 The gVisor Authors.
//
// 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 test provides functionality used to test the remote checker.
package test

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"golang.org/x/sys/unix"
	"google.golang.org/protobuf/proto"
	"gvisor.dev/gvisor/pkg/cleanup"
	"gvisor.dev/gvisor/pkg/log"
	"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/wire"
	pb "gvisor.dev/gvisor/pkg/sentry/seccheck/points/points_go_proto"
	"gvisor.dev/gvisor/pkg/sync"
	"gvisor.dev/gvisor/pkg/test/testutil"
	"gvisor.dev/gvisor/pkg/unet"
)

// Server is the counterpart to the checkers.Remote. It receives connections
// remote checkers and stores all points that it receives.
type Server struct {
	Path   string
	socket *unet.ServerSocket

	mu sync.Mutex

	// +checklocks:mu
	clients []*unet.Socket

	// +checklocks:mu
	points []Message

	version uint32
}

// Message corresponds to a single message sent from checkers.Remote.
type Message struct {
	// MsgType indicates what is the type of Msg.
	MsgType pb.MessageType
	// Msg is the payload to the message that can be decoded using MsgType.
	Msg []byte
}

// NewServer creates a new server that listens to a UDS that it creates under
// os.TempDir.
func NewServer() (*Server, error) {
	dir, err := ioutil.TempDir(os.TempDir(), "remote")
	if err != nil {
		return nil, err
	}
	server, err := newServerPath(filepath.Join(dir, "remote.sock"))
	if err != nil {
		_ = os.RemoveAll(dir)
		return nil, err
	}
	return server, nil
}

func newServerPath(path string) (*Server, error) {
	socket, err := unix.Socket(unix.AF_UNIX, unix.SOCK_SEQPACKET, 0)
	if err != nil {
		return nil, fmt.Errorf("socket(AF_UNIX, SOCK_SEQPACKET, 0): %w", err)
	}
	cu := cleanup.Make(func() {
		_ = unix.Close(socket)
	})
	defer cu.Clean()

	sa := &unix.SockaddrUnix{Name: path}
	if err := unix.Bind(socket, sa); err != nil {
		return nil, fmt.Errorf("bind(%q): %w", path, err)
	}

	ss, err := unet.NewServerSocket(socket)
	if err != nil {
		return nil, err
	}
	cu.Add(func() { ss.Close() })

	if err := ss.Listen(); err != nil {
		return nil, err
	}

	server := &Server{
		Path:    path,
		socket:  ss,
		version: wire.CurrentVersion,
	}
	go server.run()
	cu.Release()
	return server, nil
}

func (s *Server) run() {
	for {
		client, err := s.socket.Accept()
		if err != nil {
			// EBADF returns when the socket closes.
			if !errors.Is(err, unix.EBADF) {
				log.Warningf("socket.Accept(): %v", err)
			}
			return
		}
		if err := s.handshake(client); err != nil {
			log.Warningf(err.Error())
			_ = client.Close()
			continue
		}
		s.mu.Lock()
		s.clients = append(s.clients, client)
		s.mu.Unlock()
		go s.handleClient(client)
	}
}

// handshake performs version exchange with client. See common.proto for details
// about the protocol.
func (s *Server) handshake(client *unet.Socket) error {
	var in [1024]byte
	read, err := client.Read(in[:])
	if err != nil {
		return fmt.Errorf("reading handshake message: %w", err)
	}
	hsIn := pb.Handshake{}
	if err := proto.Unmarshal(in[:read], &hsIn); err != nil {
		return fmt.Errorf("unmarshalling handshake message: %w", err)
	}
	if hsIn.Version != wire.CurrentVersion {
		return fmt.Errorf("wrong version number, want: %d, got, %d", wire.CurrentVersion, hsIn.Version)
	}

	hsOut := pb.Handshake{Version: s.version}
	out, err := proto.Marshal(&hsOut)
	if err != nil {
		return fmt.Errorf("marshalling handshake message: %w", err)
	}
	if _, err := client.Write(out); err != nil {
		return fmt.Errorf("sending handshake message: %w", err)
	}
	return nil
}

func (s *Server) handleClient(client *unet.Socket) {
	defer func() {
		s.mu.Lock()
		for i, c := range s.clients {
			if c == client {
				s.clients = append(s.clients[:i], s.clients[i+1:]...)
				break
			}
		}
		s.mu.Unlock()
		_ = client.Close()
	}()

	var buf = make([]byte, 1024*1024)
	for {
		read, err := client.Read(buf)
		if err != nil {
			return
		}
		if read == 0 {
			return
		}
		if read < wire.HeaderStructSize {
			panic("invalid message")
		}
		hdr := wire.Header{}
		hdr.UnmarshalUnsafe(buf[0:wire.HeaderStructSize])
		if read < int(hdr.HeaderSize) {
			panic(fmt.Sprintf("message truncated, header size: %d, readL %d", hdr.HeaderSize, read))
		}
		msg := Message{
			MsgType: pb.MessageType(hdr.MessageType),
			Msg:     buf[hdr.HeaderSize:read],
		}
		s.mu.Lock()
		s.points = append(s.points, msg)
		s.mu.Unlock()
	}
}

// Count return the number of points it has received.
func (s *Server) Count() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.points)
}

// Reset throws aways all points received so far and returns the number of
// points discarded.
func (s *Server) Reset() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	count := len(s.points)
	s.points = nil
	return count
}

// GetPoints returns all points that it has received.
func (s *Server) GetPoints() []Message {
	s.mu.Lock()
	defer s.mu.Unlock()
	cpy := make([]Message, len(s.points))
	copy(cpy, s.points)
	return cpy
}

// Close stops listenning and closes all connections.
func (s *Server) Close() {
	_ = s.socket.Close()
	s.mu.Lock()
	for _, client := range s.clients {
		_ = client.Close()
	}
	s.mu.Unlock()
	_ = os.Remove(s.Path)
}

// WaitForCount waits for the number of points to reach the desired number for
// 5 seconds. It fails if not received in time.
func (s *Server) WaitForCount(count int) error {
	return testutil.Poll(func() error {
		if got := s.Count(); got < count {
			return fmt.Errorf("waiting for points %d to arrive, received %d", count, got)
		}
		return nil
	}, 5*time.Second)
}

// SetVersion sets the version to be used in handshake.
func (s *Server) SetVersion(newVersion uint32) {
	s.version = newVersion
}
<commit_msg>Fix data race in remote_test.<commit_after>// Copyright 2022 The gVisor Authors.
//
// 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 test provides functionality used to test the remote checker.
package test

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"golang.org/x/sys/unix"
	"google.golang.org/protobuf/proto"
	"gvisor.dev/gvisor/pkg/cleanup"
	"gvisor.dev/gvisor/pkg/log"
	"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/wire"
	pb "gvisor.dev/gvisor/pkg/sentry/seccheck/points/points_go_proto"
	"gvisor.dev/gvisor/pkg/sync"
	"gvisor.dev/gvisor/pkg/test/testutil"
	"gvisor.dev/gvisor/pkg/unet"
)

// Server is the counterpart to the checkers.Remote. It receives connections
// remote checkers and stores all points that it receives.
type Server struct {
	Path   string
	socket *unet.ServerSocket

	mu sync.Mutex

	// +checklocks:mu
	clients []*unet.Socket

	// +checklocks:mu
	points []Message

	// +checklocks:mu
	version uint32
}

// Message corresponds to a single message sent from checkers.Remote.
type Message struct {
	// MsgType indicates what is the type of Msg.
	MsgType pb.MessageType
	// Msg is the payload to the message that can be decoded using MsgType.
	Msg []byte
}

// NewServer creates a new server that listens to a UDS that it creates under
// os.TempDir.
func NewServer() (*Server, error) {
	dir, err := ioutil.TempDir(os.TempDir(), "remote")
	if err != nil {
		return nil, err
	}
	server, err := newServerPath(filepath.Join(dir, "remote.sock"))
	if err != nil {
		_ = os.RemoveAll(dir)
		return nil, err
	}
	return server, nil
}

func newServerPath(path string) (*Server, error) {
	socket, err := unix.Socket(unix.AF_UNIX, unix.SOCK_SEQPACKET, 0)
	if err != nil {
		return nil, fmt.Errorf("socket(AF_UNIX, SOCK_SEQPACKET, 0): %w", err)
	}
	cu := cleanup.Make(func() {
		_ = unix.Close(socket)
	})
	defer cu.Clean()

	sa := &unix.SockaddrUnix{Name: path}
	if err := unix.Bind(socket, sa); err != nil {
		return nil, fmt.Errorf("bind(%q): %w", path, err)
	}

	ss, err := unet.NewServerSocket(socket)
	if err != nil {
		return nil, err
	}
	cu.Add(func() { ss.Close() })

	if err := ss.Listen(); err != nil {
		return nil, err
	}

	server := &Server{
		Path:    path,
		socket:  ss,
		version: wire.CurrentVersion,
	}
	go server.run()
	cu.Release()
	return server, nil
}

func (s *Server) run() {
	for {
		client, err := s.socket.Accept()
		if err != nil {
			// EBADF returns when the socket closes.
			if !errors.Is(err, unix.EBADF) {
				log.Warningf("socket.Accept(): %v", err)
			}
			return
		}
		if err := s.handshake(client); err != nil {
			log.Warningf(err.Error())
			_ = client.Close()
			continue
		}
		s.mu.Lock()
		s.clients = append(s.clients, client)
		s.mu.Unlock()
		go s.handleClient(client)
	}
}

// handshake performs version exchange with client. See common.proto for details
// about the protocol.
func (s *Server) handshake(client *unet.Socket) error {
	var in [1024]byte
	read, err := client.Read(in[:])
	if err != nil {
		return fmt.Errorf("reading handshake message: %w", err)
	}
	hsIn := pb.Handshake{}
	if err := proto.Unmarshal(in[:read], &hsIn); err != nil {
		return fmt.Errorf("unmarshalling handshake message: %w", err)
	}
	if hsIn.Version != wire.CurrentVersion {
		return fmt.Errorf("wrong version number, want: %d, got, %d", wire.CurrentVersion, hsIn.Version)
	}

	s.mu.Lock()
	v := s.version
	s.mu.Unlock()
	hsOut := pb.Handshake{Version: v}
	out, err := proto.Marshal(&hsOut)
	if err != nil {
		return fmt.Errorf("marshalling handshake message: %w", err)
	}
	if _, err := client.Write(out); err != nil {
		return fmt.Errorf("sending handshake message: %w", err)
	}
	return nil
}

func (s *Server) handleClient(client *unet.Socket) {
	defer func() {
		s.mu.Lock()
		for i, c := range s.clients {
			if c == client {
				s.clients = append(s.clients[:i], s.clients[i+1:]...)
				break
			}
		}
		s.mu.Unlock()
		_ = client.Close()
	}()

	var buf = make([]byte, 1024*1024)
	for {
		read, err := client.Read(buf)
		if err != nil {
			return
		}
		if read == 0 {
			return
		}
		if read < wire.HeaderStructSize {
			panic("invalid message")
		}
		hdr := wire.Header{}
		hdr.UnmarshalUnsafe(buf[0:wire.HeaderStructSize])
		if read < int(hdr.HeaderSize) {
			panic(fmt.Sprintf("message truncated, header size: %d, readL %d", hdr.HeaderSize, read))
		}
		msg := Message{
			MsgType: pb.MessageType(hdr.MessageType),
			Msg:     buf[hdr.HeaderSize:read],
		}
		s.mu.Lock()
		s.points = append(s.points, msg)
		s.mu.Unlock()
	}
}

// Count return the number of points it has received.
func (s *Server) Count() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.points)
}

// Reset throws aways all points received so far and returns the number of
// points discarded.
func (s *Server) Reset() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	count := len(s.points)
	s.points = nil
	return count
}

// GetPoints returns all points that it has received.
func (s *Server) GetPoints() []Message {
	s.mu.Lock()
	defer s.mu.Unlock()
	cpy := make([]Message, len(s.points))
	copy(cpy, s.points)
	return cpy
}

// Close stops listenning and closes all connections.
func (s *Server) Close() {
	_ = s.socket.Close()
	s.mu.Lock()
	for _, client := range s.clients {
		_ = client.Close()
	}
	s.mu.Unlock()
	_ = os.Remove(s.Path)
}

// WaitForCount waits for the number of points to reach the desired number for
// 5 seconds. It fails if not received in time.
func (s *Server) WaitForCount(count int) error {
	return testutil.Poll(func() error {
		if got := s.Count(); got < count {
			return fmt.Errorf("waiting for points %d to arrive, received %d", count, got)
		}
		return nil
	}, 5*time.Second)
}

// SetVersion sets the version to be used in handshake.
func (s *Server) SetVersion(newVersion uint32) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.version = newVersion
}
<|endoftext|>


<commit_before>package dmgo

import "fmt"

type mem struct {
	cart    []byte
	cartRAM []byte

	internalRAM     [0x8000]byte // go ahead and do CGB size
	highInternalRAM [0x7f]byte   // go ahead and do CGB size
	mbc             mbc
}

func (mem *mem) mbcRead(addr uint16) byte {
	return mem.mbc.Read(mem, addr)
}
func (mem *mem) mbcWrite(addr uint16, val byte) {
	mem.mbc.Write(mem, addr, val)
}

// TODO / FIXME: DMA timing enforcement!
// also, dma cancel / restart (if you
// run a 2nd dma, it cancels the running one
// and starts the new one immediately)
func (cs *cpuState) oamDMA(addr uint16) {
	for i := uint16(0); i < 0xa0; i++ {
		cs.write(0xfe00+i, cs.read(addr+i))
	}
}

func (cs *cpuState) read(addr uint16) byte {
	var val byte
	switch {

	case addr < 0x8000:
		val = cs.mem.mbcRead(addr)

	case addr >= 0x8000 && addr < 0xa000:
		val = cs.lcd.videoRAM[addr-0x8000]

	case addr >= 0xa000 && addr < 0xc000:
		val = cs.mem.mbcRead(addr)

	case addr >= 0xc000 && addr < 0xfe00:
		ramAddr := (addr - 0xc000) & 0x1fff // 8kb with wraparound
		val = cs.mem.internalRAM[ramAddr]

	case addr == 0xff00:
		val = cs.joypad.readJoypadReg()
	case addr == 0xff02:
		val = cs.readSerialControlReg()

	case addr == 0xff04:
		val = byte(cs.timerDivCycles >> 8)
	case addr == 0xff05:
		val = cs.timerCounterReg
	case addr == 0xff06:
		val = cs.timerModuloReg
	case addr == 0xff07:
		val = cs.readTimerControlReg()

	case addr == 0xff0f:
		val = cs.readInterruptFlagReg()

	case addr == 0xff12:
		val = cs.apu.sounds[0].readSoundEnvReg()
	case addr == 0xff14:
		val = cs.apu.sounds[0].readFreqHighReg()
	case addr == 0xff17:
		val = cs.apu.sounds[1].readSoundEnvReg()
	case addr == 0xff19:
		val = cs.apu.sounds[1].readFreqHighReg()
	case addr == 0xff1c:
		val = cs.apu.sounds[2].readWaveOutLvlReg()
	case addr == 0xff1e:
		val = cs.apu.sounds[2].readFreqHighReg()
	case addr == 0xff21:
		val = cs.apu.sounds[3].readSoundEnvReg()
	case addr == 0xff23:
		val = cs.apu.sounds[3].readFreqHighReg()
	case addr == 0xff24:
		val = cs.apu.readVolumeReg()
	case addr == 0xff25:
		val = cs.apu.readSpeakerSelectReg()
	case addr == 0xff26:
		val = cs.apu.readSoundOnOffReg()

	case addr >= 0xff30 && addr < 0xff40:
		val = cs.apu.sounds[2].wavePatternRAM[addr-0xff30]

	case addr == 0xff40:
		val = cs.lcd.readControlReg()
	case addr == 0xff41:
		val = cs.lcd.readStatusReg()
	case addr == 0xff42:
		val = cs.lcd.scrollY
	case addr == 0xff43:
		val = cs.lcd.scrollX
	case addr == 0xff44:
		val = cs.lcd.lyReg
	case addr == 0xff45:
		val = cs.lcd.lycReg
	case addr == 0xff47:
		val = cs.lcd.backgroundPaletteReg
	case addr == 0xff48:
		val = cs.lcd.objectPalette0Reg
	case addr == 0xff49:
		val = cs.lcd.objectPalette1Reg
	case addr == 0xff4a:
		val = cs.lcd.windowY
	case addr == 0xff4b:
		val = cs.lcd.windowX

	case addr >= 0xff80 && addr < 0xffff:
		val = cs.mem.highInternalRAM[addr-0xff80]
	case addr == 0xffff:
		val = cs.readInterruptEnableReg()

	default:
		cs.stepErr(fmt.Sprintf("not implemented: read at %x\n", addr))
	}
	//	fmt.Printf("\treading 0x%02x from 0x%04x\n", val, addr)
	return val
}

func (cs *cpuState) read16(addr uint16) uint16 {
	high := uint16(cs.read(addr + 1))
	low := uint16(cs.read(addr))
	return (high << 8) | low
}

func (cs *cpuState) write(addr uint16, val byte) {
	switch {

	case addr < 0x8000:
		cs.mem.mbcWrite(addr, val)

	case addr >= 0x8000 && addr < 0xa000:
		cs.lcd.writeVideoRAM(addr-0x8000, val)

	case addr >= 0xa000 && addr < 0xc000:
		cs.mem.mbcWrite(addr, val)

	case addr >= 0xc000 && addr < 0xfe00:
		cs.mem.internalRAM[((addr - 0xc000) & 0x1fff)] = val // 8kb with wraparound
	case addr >= 0xfe00 && addr < 0xfea0:
		cs.lcd.writeOAM(addr-0xfe00, val)
	case addr >= 0xfea0 && addr < 0xff00:
		// empty, nop (can be more complicated, see TCAGBD)

	case addr == 0xff00:
		cs.joypad.writeJoypadReg(val)
	case addr == 0xff01:
		cs.serialTransferData = val
	case addr == 0xff02:
		cs.writeSerialControlReg(val)

	case addr == 0xff04:
		cs.timerDivCycles = 0
	case addr == 0xff05:
		cs.timerCounterReg = val
	case addr == 0xff06:
		cs.timerModuloReg = val
	case addr == 0xff07:
		cs.writeTimerControlReg(val)

	case addr == 0xff10:
		cs.apu.sounds[0].writeSweepReg(val)
	case addr == 0xff11:
		cs.apu.sounds[0].writeLenDutyReg(val)
	case addr == 0xff12:
		cs.apu.sounds[0].writeSoundEnvReg(val)
	case addr == 0xff13:
		cs.apu.sounds[0].writeFreqLowReg(val)
	case addr == 0xff14:
		cs.apu.sounds[0].writeFreqHighReg(val)

	case addr == 0xff16:
		cs.apu.sounds[1].writeLenDutyReg(val)
	case addr == 0xff17:
		cs.apu.sounds[1].writeSoundEnvReg(val)
	case addr == 0xff18:
		cs.apu.sounds[1].writeFreqLowReg(val)
	case addr == 0xff19:
		cs.apu.sounds[1].writeFreqHighReg(val)

	case addr == 0xff1a:
		cs.apu.sounds[2].on = val&0x80 != 0
	case addr == 0xff1b:
		cs.apu.sounds[2].lengthData = val
	case addr == 0xff1c:
		cs.apu.sounds[2].writeWaveOutLvlReg(val)
	case addr == 0xff1d:
		cs.apu.sounds[2].writeFreqLowReg(val)
	case addr == 0xff1e:
		cs.apu.sounds[2].writeFreqHighReg(val)

	case addr == 0xff20:
		cs.apu.sounds[3].lengthData = val & 0x1f
	case addr == 0xff21:
		cs.apu.sounds[3].writeSoundEnvReg(val)
	case addr == 0xff22:
		cs.apu.sounds[3].writePolyCounterReg(val)
	case addr == 0xff23:
		cs.apu.sounds[3].writeFreqHighReg(val) // noise channel uses control bits, freq ignored

	case addr == 0xff24:
		cs.apu.writeVolumeReg(val)
	case addr == 0xff25:
		cs.apu.writeSpeakerSelectReg(val)
	case addr == 0xff26:
		cs.apu.writeSoundOnOffReg(val)
	case addr >= 0xff30 && addr < 0xff40:
		cs.apu.sounds[2].wavePatternRAM[addr-0xff30] = val

	case addr == 0xff40:
		cs.lcd.writeControlReg(val)
	case addr == 0xff41:
		cs.lcd.writeStatusReg(val)
	case addr == 0xff42:
		cs.lcd.writeScrollY(val)
	case addr == 0xff43:
		cs.lcd.writeScrollX(val)
	case addr == 0xff44:
		// nop? pandocs says something
		// about "resetting the counter",
		// bgb seems to do nothing. doesn't
		// reset lyReg, doesn't change any
		// counter I see...
	case addr == 0xff45:
		cs.lcd.writeLycReg(val)
	case addr == 0xff46:
		cs.oamDMA(uint16(val) << 8)
	case addr == 0xff47:
		cs.lcd.writeBackgroundPaletteReg(val)
	case addr == 0xff48:
		cs.lcd.writeObjectPalette0Reg(val)
	case addr == 0xff49:
		cs.lcd.writeObjectPalette1Reg(val)
	case addr == 0xff4a:
		cs.lcd.writeWindowY(val)
	case addr == 0xff4b:
		cs.lcd.writeWindowX(val)

	case addr == 0xff0f:
		cs.writeInterruptFlagReg(val)
	case addr >= 0xff4c && addr < 0xff80:
		// empty, nop (can be more complicated, see TCAGBD)
	case addr >= 0xff80 && addr < 0xffff:
		cs.mem.highInternalRAM[addr-0xff80] = val
	case addr == 0xffff:
		cs.writeInterruptEnableReg(val)
	default:
		cs.stepErr(fmt.Sprintf("not implemented: write(0x%04x, %v)\n", addr, val))
	}
	//	fmt.Printf("\twriting 0x%02x to 0x%04x\n", val, addr)
}

func (cs *cpuState) write16(addr uint16, val uint16) {
	cs.write(addr, byte(val))
	cs.write(addr+1, byte(val>>8))
}
<commit_msg>More mem map, fix case order bug (go has those)<commit_after>package dmgo

import "fmt"

type mem struct {
	cart    []byte
	cartRAM []byte

	internalRAM     [0x8000]byte // go ahead and do CGB size
	highInternalRAM [0x7f]byte   // go ahead and do CGB size
	mbc             mbc
}

func (mem *mem) mbcRead(addr uint16) byte {
	return mem.mbc.Read(mem, addr)
}
func (mem *mem) mbcWrite(addr uint16, val byte) {
	mem.mbc.Write(mem, addr, val)
}

// TODO / FIXME: DMA timing enforcement!
// also, dma cancel / restart (if you
// run a 2nd dma, it cancels the running one
// and starts the new one immediately)
func (cs *cpuState) oamDMA(addr uint16) {
	for i := uint16(0); i < 0xa0; i++ {
		cs.write(0xfe00+i, cs.read(addr+i))
	}
}

func (cs *cpuState) read(addr uint16) byte {
	var val byte
	switch {

	case addr < 0x8000:
		val = cs.mem.mbcRead(addr)

	case addr >= 0x8000 && addr < 0xa000:
		val = cs.lcd.videoRAM[addr-0x8000]

	case addr >= 0xa000 && addr < 0xc000:
		val = cs.mem.mbcRead(addr)

	case addr >= 0xc000 && addr < 0xfe00:
		ramAddr := (addr - 0xc000) & 0x1fff // 8kb with wraparound
		val = cs.mem.internalRAM[ramAddr]

	case addr == 0xff00:
		val = cs.joypad.readJoypadReg()
	case addr == 0xff02:
		val = cs.readSerialControlReg()

	case addr == 0xff04:
		val = byte(cs.timerDivCycles >> 8)
	case addr == 0xff05:
		val = cs.timerCounterReg
	case addr == 0xff06:
		val = cs.timerModuloReg
	case addr == 0xff07:
		val = cs.readTimerControlReg()

	case addr == 0xff0f:
		val = cs.readInterruptFlagReg()

	case addr == 0xff11:
		val = cs.apu.sounds[0].readLenDutyReg()
	case addr == 0xff12:
		val = cs.apu.sounds[0].readSoundEnvReg()
	case addr == 0xff14:
		val = cs.apu.sounds[0].readFreqHighReg()
	case addr == 0xff16:
		val = cs.apu.sounds[1].readLenDutyReg()
	case addr == 0xff17:
		val = cs.apu.sounds[1].readSoundEnvReg()
	case addr == 0xff19:
		val = cs.apu.sounds[1].readFreqHighReg()
	case addr == 0xff1b:
		val = cs.apu.sounds[2].lengthData
	case addr == 0xff1c:
		val = cs.apu.sounds[2].readWaveOutLvlReg()
	case addr == 0xff1e:
		val = cs.apu.sounds[2].readFreqHighReg()
	case addr == 0xff21:
		val = cs.apu.sounds[3].readSoundEnvReg()
	case addr == 0xff23:
		val = cs.apu.sounds[3].readFreqHighReg()
	case addr == 0xff24:
		val = cs.apu.readVolumeReg()
	case addr == 0xff25:
		val = cs.apu.readSpeakerSelectReg()
	case addr == 0xff26:
		val = cs.apu.readSoundOnOffReg()

	case addr >= 0xff30 && addr < 0xff40:
		val = cs.apu.sounds[2].wavePatternRAM[addr-0xff30]

	case addr == 0xff40:
		val = cs.lcd.readControlReg()
	case addr == 0xff41:
		val = cs.lcd.readStatusReg()
	case addr == 0xff42:
		val = cs.lcd.scrollY
	case addr == 0xff43:
		val = cs.lcd.scrollX
	case addr == 0xff44:
		val = cs.lcd.lyReg
	case addr == 0xff45:
		val = cs.lcd.lycReg
	case addr == 0xff47:
		val = cs.lcd.backgroundPaletteReg
	case addr == 0xff48:
		val = cs.lcd.objectPalette0Reg
	case addr == 0xff49:
		val = cs.lcd.objectPalette1Reg
	case addr == 0xff4a:
		val = cs.lcd.windowY
	case addr == 0xff4b:
		val = cs.lcd.windowX

	case addr >= 0xff4c && addr < 0xff80:
		val = 0xff // unmapped io bytes

	case addr >= 0xff80 && addr < 0xffff:
		val = cs.mem.highInternalRAM[addr-0xff80]
	case addr == 0xffff:
		val = cs.readInterruptEnableReg()

	default:
		cs.stepErr(fmt.Sprintf("not implemented: read at %x\n", addr))
	}
	//	fmt.Printf("\treading 0x%02x from 0x%04x\n", val, addr)
	return val
}

func (cs *cpuState) read16(addr uint16) uint16 {
	high := uint16(cs.read(addr + 1))
	low := uint16(cs.read(addr))
	return (high << 8) | low
}

func (cs *cpuState) write(addr uint16, val byte) {
	switch {

	case addr < 0x8000:
		cs.mem.mbcWrite(addr, val)

	case addr >= 0x8000 && addr < 0xa000:
		cs.lcd.writeVideoRAM(addr-0x8000, val)

	case addr >= 0xa000 && addr < 0xc000:
		cs.mem.mbcWrite(addr, val)

	case addr >= 0xc000 && addr < 0xfe00:
		cs.mem.internalRAM[((addr - 0xc000) & 0x1fff)] = val // 8kb with wraparound
	case addr >= 0xfe00 && addr < 0xfea0:
		cs.lcd.writeOAM(addr-0xfe00, val)
	case addr >= 0xfea0 && addr < 0xff00:
		// empty, nop (can be more complicated, see TCAGBD)

	case addr == 0xff00:
		cs.joypad.writeJoypadReg(val)
	case addr == 0xff01:
		cs.serialTransferData = val
	case addr == 0xff02:
		cs.writeSerialControlReg(val)

	case addr == 0xff04:
		cs.timerDivCycles = 0
	case addr == 0xff05:
		cs.timerCounterReg = val
	case addr == 0xff06:
		cs.timerModuloReg = val
	case addr == 0xff07:
		cs.writeTimerControlReg(val)
	case addr == 0xff0f:
		cs.writeInterruptFlagReg(val)

	case addr == 0xff10:
		cs.apu.sounds[0].writeSweepReg(val)
	case addr == 0xff11:
		cs.apu.sounds[0].writeLenDutyReg(val)
	case addr == 0xff12:
		cs.apu.sounds[0].writeSoundEnvReg(val)
	case addr == 0xff13:
		cs.apu.sounds[0].writeFreqLowReg(val)
	case addr == 0xff14:
		cs.apu.sounds[0].writeFreqHighReg(val)

	case addr == 0xff16:
		cs.apu.sounds[1].writeLenDutyReg(val)
	case addr == 0xff17:
		cs.apu.sounds[1].writeSoundEnvReg(val)
	case addr == 0xff18:
		cs.apu.sounds[1].writeFreqLowReg(val)
	case addr == 0xff19:
		cs.apu.sounds[1].writeFreqHighReg(val)

	case addr == 0xff1a:
		cs.apu.sounds[2].on = val&0x80 != 0
	case addr == 0xff1b:
		cs.apu.sounds[2].lengthData = val
	case addr == 0xff1c:
		cs.apu.sounds[2].writeWaveOutLvlReg(val)
	case addr == 0xff1d:
		cs.apu.sounds[2].writeFreqLowReg(val)
	case addr == 0xff1e:
		cs.apu.sounds[2].writeFreqHighReg(val)

	case addr == 0xff20:
		cs.apu.sounds[3].lengthData = val & 0x1f
	case addr == 0xff21:
		cs.apu.sounds[3].writeSoundEnvReg(val)
	case addr == 0xff22:
		cs.apu.sounds[3].writePolyCounterReg(val)
	case addr == 0xff23:
		cs.apu.sounds[3].writeFreqHighReg(val) // noise channel uses control bits, freq ignored

	case addr == 0xff24:
		cs.apu.writeVolumeReg(val)
	case addr == 0xff25:
		cs.apu.writeSpeakerSelectReg(val)
	case addr == 0xff26:
		cs.apu.writeSoundOnOffReg(val)
	case addr >= 0xff30 && addr < 0xff40:
		cs.apu.sounds[2].wavePatternRAM[addr-0xff30] = val

	case addr == 0xff40:
		cs.lcd.writeControlReg(val)
	case addr == 0xff41:
		cs.lcd.writeStatusReg(val)
	case addr == 0xff42:
		cs.lcd.writeScrollY(val)
	case addr == 0xff43:
		cs.lcd.writeScrollX(val)
	case addr == 0xff44:
		// nop? pandocs says something
		// about "resetting the counter",
		// bgb seems to do nothing. doesn't
		// reset lyReg, doesn't change any
		// counter I see...
	case addr == 0xff45:
		cs.lcd.writeLycReg(val)
	case addr == 0xff46:
		cs.oamDMA(uint16(val) << 8)
	case addr == 0xff47:
		cs.lcd.writeBackgroundPaletteReg(val)
	case addr == 0xff48:
		cs.lcd.writeObjectPalette0Reg(val)
	case addr == 0xff49:
		cs.lcd.writeObjectPalette1Reg(val)
	case addr == 0xff4a:
		cs.lcd.writeWindowY(val)
	case addr == 0xff4b:
		cs.lcd.writeWindowX(val)

	case addr >= 0xff4c && addr < 0xff80:
		// empty, nop (can be more complicated, see TCAGBD)
	case addr >= 0xff80 && addr < 0xffff:
		cs.mem.highInternalRAM[addr-0xff80] = val
	case addr == 0xffff:
		cs.writeInterruptEnableReg(val)
	default:
		cs.stepErr(fmt.Sprintf("not implemented: write(0x%04x, %v)\n", addr, val))
	}
}

func (cs *cpuState) write16(addr uint16, val uint16) {
	cs.write(addr, byte(val))
	cs.write(addr+1, byte(val>>8))
}
<|endoftext|>


<commit_before>package rpc

import (
	"code.google.com/p/go.net/websocket"
	"errors"
	"koding/newkite/dnode"
	"net/url"
	"time"
)

const redialDurationStart = 1 * time.Second
const redialDurationMax = 60 * time.Second

// Dial is a helper for creating a Client for just calling methods on the server.
// Do not use it if you want to handle methods on client side. Instead create a
// new Client, register your methods on Client.Dnode then call Client.Dial().
func Dial(url string, reconnect bool) (*Client, error) {
	c := NewClient()
	c.Reconnect = reconnect

	err := c.Dial(url)
	if err != nil {
		return nil, err
	}

	return c, nil
}

// Client is a dnode RPC client.
type Client struct {
	// Websocket connection
	Conn *websocket.Conn

	// Websocket connection options.
	Config *websocket.Config

	// Dnode message processor.
	dnode *dnode.Dnode

	// A space for saving/reading extra properties about this client.
	properties map[string]interface{}

	// Should we reconnect if disconnected?
	Reconnect bool

	// Time to wait before redial connection.
	redialDuration time.Duration

	// on connect/disconnect handlers are invoked after every
	// connect/disconnect.
	onConnectHandlers    []func()
	onDisconnectHandlers []func()
}

// NewClient returns a pointer to new Client.
// You need to call Dial() before interacting with the Server.
func NewClient() *Client {
	// Must send an "Origin" header. Does not checked on server.
	origin, _ := url.Parse("")

	config := &websocket.Config{
		Version: websocket.ProtocolVersionHybi13,
		Origin:  origin,
		// Location will be set when dialing.
	}

	c := &Client{
		properties:     make(map[string]interface{}),
		redialDuration: redialDurationStart,
		Config:         config,
	}

	c.dnode = dnode.New(c)
	return c
}

func (c *Client) SetWrappers(wrapMethodArgs, wrapCallbackArgs dnode.Wrapper, runMethod, runCallback dnode.Runner) {
	c.dnode.WrapMethodArgs = wrapMethodArgs
	c.dnode.WrapCallbackArgs = wrapCallbackArgs
	c.dnode.RunMethod = runMethod
	c.dnode.RunCallback = runCallback
}

// Dial connects to the dnode server on "url" and starts a goroutine
// that processes incoming messages.
//
// Do not forget to register your handlers on Client.Dnode
// before calling Dial() to prevent race conditions.
func (c *Client) Dial(url_ string) (err error) {
	if c.Config.Location, err = url.Parse(url_); err != nil {
		return
	}

	if err = c.dial(); err != nil {
		return
	}

	go c.run()

	return
}

// dial makes a single Dial() and run onConnectHandlers if connects.
func (c *Client) dial() (err error) {
	ws, err := websocket.DialConfig(c.Config)
	if err != nil {
		return err
	}

	// We are connected
	c.Conn = ws

	// Reset the wait time.
	c.redialDuration = redialDurationStart

	// Must be run in a goroutine because a handler may wait a response from
	// server.
	go c.callOnConnectHandlers()

	return nil
}

// DialForever connects to the server in background.
// If the connection drops, it reconnects again.
func (c *Client) DialForever(url_ string) (err error) {
	if c.Config.Location, err = url.Parse(url_); err != nil {
		return
	}

	go c.dialForever()

	return
}

func (c *Client) dialForever() {
	for c.dial() != nil {
		if !c.Reconnect {
			return
		}

		c.sleep()
	}
	go c.run()
}

// run consumes incoming dnode messages. Reconnects if necessary.
func (c *Client) run() (err error) {
	for {
	running:
		err = c.dnode.Run()
		c.callOnDisconnectHandlers()
	dialAgain:
		if !c.Reconnect {
			break
		}

		err = c.dial()
		if err != nil {
			c.sleep()
			goto dialAgain
		}

		goto running
	}

	return err
}

// sleep is used to wait for a while between dial retries.
// Each time it is called the redialDuration is incremented.
func (c *Client) sleep() {
	time.Sleep(c.redialDuration)

	c.redialDuration *= 2
	if c.redialDuration > redialDurationMax {
		c.redialDuration = redialDurationMax
	}
}

// Close closes the underlying websocket connection.
func (c *Client) Close() {
	c.Conn.Close()
}

func (c *Client) Send(msg []byte) error {
	// println("\nSending...", string(msg))
	if c.Conn == nil {
		return errors.New("Not connected")
	}

	return websocket.Message.Send(c.Conn, string(msg))
}

func (c *Client) Receive() ([]byte, error) {
	// println("Receiving...")
	var msg []byte
	err := websocket.Message.Receive(c.Conn, &msg)
	// println("\nReceived:", string(msg))
	return msg, err
}

func (c *Client) RemoveCallback(id uint64) {
	c.dnode.RemoveCallback(id)
}

// RemoteAddr returns the host:port as string if server connection.
func (c *Client) RemoteAddr() string {
	if c.Conn.IsServerConn() {
		return c.Conn.Request().RemoteAddr
	}
	return ""
}

func (c *Client) Properties() map[string]interface{} {
	return c.properties
}

// Call calls a method with args on the dnode server.
func (c *Client) Call(method string, args ...interface{}) (map[string]dnode.Path, error) {
	return c.dnode.Call(method, args...)
}

// OnConnect registers a function to run on client connect.
func (c *Client) OnConnect(handler func()) {
	c.onConnectHandlers = append(c.onConnectHandlers, handler)
}

// OnDisconnect registers a function to run on client disconnect.
func (c *Client) OnDisconnect(handler func()) {
	c.onDisconnectHandlers = append(c.onDisconnectHandlers, handler)
}

// callOnConnectHandlers runs the registered connect handlers.
func (c *Client) callOnConnectHandlers() {
	for _, handler := range c.onConnectHandlers {
		go handler()
	}
}

// callOnDisconnectHandlers runs the registered disconnect handlers.
func (c *Client) callOnDisconnectHandlers() {
	for _, handler := range c.onDisconnectHandlers {
		go handler()
	}
}
<commit_msg>kite: remove named results<commit_after>package rpc

import (
	"code.google.com/p/go.net/websocket"
	"errors"
	"koding/newkite/dnode"
	"net/url"
	"time"
)

const redialDurationStart = 1 * time.Second
const redialDurationMax = 60 * time.Second

// Dial is a helper for creating a Client for just calling methods on the server.
// Do not use it if you want to handle methods on client side. Instead create a
// new Client, register your methods on Client.Dnode then call Client.Dial().
func Dial(url string, reconnect bool) (*Client, error) {
	c := NewClient()
	c.Reconnect = reconnect

	err := c.Dial(url)
	if err != nil {
		return nil, err
	}

	return c, nil
}

// Client is a dnode RPC client.
type Client struct {
	// Websocket connection
	Conn *websocket.Conn

	// Websocket connection options.
	Config *websocket.Config

	// Dnode message processor.
	dnode *dnode.Dnode

	// A space for saving/reading extra properties about this client.
	properties map[string]interface{}

	// Should we reconnect if disconnected?
	Reconnect bool

	// Time to wait before redial connection.
	redialDuration time.Duration

	// on connect/disconnect handlers are invoked after every
	// connect/disconnect.
	onConnectHandlers    []func()
	onDisconnectHandlers []func()
}

// NewClient returns a pointer to new Client.
// You need to call Dial() before interacting with the Server.
func NewClient() *Client {
	// Must send an "Origin" header. Does not checked on server.
	origin, _ := url.Parse("")

	config := &websocket.Config{
		Version: websocket.ProtocolVersionHybi13,
		Origin:  origin,
		// Location will be set when dialing.
	}

	c := &Client{
		properties:     make(map[string]interface{}),
		redialDuration: redialDurationStart,
		Config:         config,
	}

	c.dnode = dnode.New(c)
	return c
}

func (c *Client) SetWrappers(wrapMethodArgs, wrapCallbackArgs dnode.Wrapper, runMethod, runCallback dnode.Runner) {
	c.dnode.WrapMethodArgs = wrapMethodArgs
	c.dnode.WrapCallbackArgs = wrapCallbackArgs
	c.dnode.RunMethod = runMethod
	c.dnode.RunCallback = runCallback
}

// Dial connects to the dnode server on "url" and starts a goroutine
// that processes incoming messages.
//
// Do not forget to register your handlers on Client.Dnode
// before calling Dial() to prevent race conditions.
func (c *Client) Dial(url_ string) error {
	var err error

	if c.Config.Location, err = url.Parse(url_); err != nil {
		return err
	}

	if err = c.dial(); err != nil {
		return err
	}

	go c.run()

	return nil
}

// dial makes a single Dial() and run onConnectHandlers if connects.
func (c *Client) dial() error {
	ws, err := websocket.DialConfig(c.Config)
	if err != nil {
		return err
	}

	// We are connected
	c.Conn = ws

	// Reset the wait time.
	c.redialDuration = redialDurationStart

	// Must be run in a goroutine because a handler may wait a response from
	// server.
	go c.callOnConnectHandlers()

	return nil
}

// DialForever connects to the server in background.
// If the connection drops, it reconnects again.
func (c *Client) DialForever(url_ string) (err error) {
	if c.Config.Location, err = url.Parse(url_); err != nil {
		return
	}

	go c.dialForever()

	return
}

func (c *Client) dialForever() {
	for c.dial() != nil {
		if !c.Reconnect {
			return
		}

		c.sleep()
	}
	go c.run()
}

// run consumes incoming dnode messages. Reconnects if necessary.
func (c *Client) run() (err error) {
	for {
	running:
		err = c.dnode.Run()
		c.callOnDisconnectHandlers()
	dialAgain:
		if !c.Reconnect {
			break
		}

		err = c.dial()
		if err != nil {
			c.sleep()
			goto dialAgain
		}

		goto running
	}

	return err
}

// sleep is used to wait for a while between dial retries.
// Each time it is called the redialDuration is incremented.
func (c *Client) sleep() {
	time.Sleep(c.redialDuration)

	c.redialDuration *= 2
	if c.redialDuration > redialDurationMax {
		c.redialDuration = redialDurationMax
	}
}

// Close closes the underlying websocket connection.
func (c *Client) Close() {
	c.Conn.Close()
}

func (c *Client) Send(msg []byte) error {
	// println("\nSending...", string(msg))
	if c.Conn == nil {
		return errors.New("Not connected")
	}

	return websocket.Message.Send(c.Conn, string(msg))
}

func (c *Client) Receive() ([]byte, error) {
	// println("Receiving...")
	var msg []byte
	err := websocket.Message.Receive(c.Conn, &msg)
	// println("\nReceived:", string(msg))
	return msg, err
}

func (c *Client) RemoveCallback(id uint64) {
	c.dnode.RemoveCallback(id)
}

// RemoteAddr returns the host:port as string if server connection.
func (c *Client) RemoteAddr() string {
	if c.Conn.IsServerConn() {
		return c.Conn.Request().RemoteAddr
	}
	return ""
}

func (c *Client) Properties() map[string]interface{} {
	return c.properties
}

// Call calls a method with args on the dnode server.
func (c *Client) Call(method string, args ...interface{}) (map[string]dnode.Path, error) {
	return c.dnode.Call(method, args...)
}

// OnConnect registers a function to run on client connect.
func (c *Client) OnConnect(handler func()) {
	c.onConnectHandlers = append(c.onConnectHandlers, handler)
}

// OnDisconnect registers a function to run on client disconnect.
func (c *Client) OnDisconnect(handler func()) {
	c.onDisconnectHandlers = append(c.onDisconnectHandlers, handler)
}

// callOnConnectHandlers runs the registered connect handlers.
func (c *Client) callOnConnectHandlers() {
	for _, handler := range c.onConnectHandlers {
		go handler()
	}
}

// callOnDisconnectHandlers runs the registered disconnect handlers.
func (c *Client) callOnDisconnectHandlers() {
	for _, handler := range c.onDisconnectHandlers {
		go handler()
	}
}
<|endoftext|>