<commit_before>package command

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/helper"
	"github.com/hashicorp/nomad/helper/flatmap"
	"github.com/kr/pretty"
	"github.com/mitchellh/cli"
)

func TestHelpers_FormatKV(t *testing.T) {
	t.Parallel()
	in := []string{"alpha|beta", "charlie|delta", "echo|"}
	out := formatKV(in)

	expect := "alpha   = beta\n"
	expect += "charlie = delta\n"
	expect += "echo    = <none>"

	if out != expect {
		t.Fatalf("expect: %s, got: %s", expect, out)
	}
}

func TestHelpers_FormatList(t *testing.T) {
	t.Parallel()
	in := []string{"alpha|beta||delta"}
	out := formatList(in)

	expect := "alpha  beta  <none>  delta"

	if out != expect {
		t.Fatalf("expect: %s, got: %s", expect, out)
	}
}

func TestHelpers_NodeID(t *testing.T) {
	t.Parallel()
	srv, _, _ := testServer(t, false, nil)
	defer srv.Shutdown()

	meta := Meta{Ui: new(cli.MockUi)}
	client, err := meta.Client()
	if err != nil {
		t.FailNow()
	}

	// This is because there is no client
	if _, err := getLocalNodeID(client); err == nil {
		t.Fatalf("getLocalNodeID() should fail")
	}
}

func TestHelpers_LineLimitReader_NoTimeLimit(t *testing.T) {
	t.Parallel()
	helloString := `hello
world
this
is
a
test`

	noLines := "jskdfhjasdhfjkajkldsfdlsjkahfkjdsafa"

	cases := []struct {
		Input       string
		Output      string
		Lines       int
		SearchLimit int
	}{
		{
			Input:       helloString,
			Output:      helloString,
			Lines:       6,
			SearchLimit: 1000,
		},
		{
			Input: helloString,
			Output: `world
this
is
a
test`,
			Lines:       5,
			SearchLimit: 1000,
		},
		{
			Input:       helloString,
			Output:      `test`,
			Lines:       1,
			SearchLimit: 1000,
		},
		{
			Input:       helloString,
			Output:      "",
			Lines:       0,
			SearchLimit: 1000,
		},
		{
			Input:       helloString,
			Output:      helloString,
			Lines:       6,
			SearchLimit: 1, // Exceed the limit
		},
		{
			Input:       noLines,
			Output:      noLines,
			Lines:       10,
			SearchLimit: 1000,
		},
		{
			Input:       noLines,
			Output:      noLines,
			Lines:       10,
			SearchLimit: 2,
		},
	}

	for i, c := range cases {
		in := ioutil.NopCloser(strings.NewReader(c.Input))
		limit := NewLineLimitReader(in, c.Lines, c.SearchLimit, 0)
		outBytes, err := ioutil.ReadAll(limit)
		if err != nil {
			t.Fatalf("case %d failed: %v", i, err)
		}

		out := string(outBytes)
		if out != c.Output {
			t.Fatalf("case %d: got %q; want %q", i, out, c.Output)
		}
	}
}

type testReadCloser struct {
	data chan []byte
}

func (t *testReadCloser) Read(p []byte) (n int, err error) {
	select {
	case b, ok := <-t.data:
		if !ok {
			return 0, io.EOF
		}

		return copy(p, b), nil
	case <-time.After(10 * time.Millisecond):
		return 0, nil
	}
}

func (t *testReadCloser) Close() error {
	close(t.data)
	return nil
}

func TestHelpers_LineLimitReader_TimeLimit(t *testing.T) {
	t.Parallel()
	// Create the test reader
	in := &testReadCloser{data: make(chan []byte)}

	// Set up the reader such that it won't hit the line/buffer limit and could
	// only terminate if it hits the time limit
	limit := NewLineLimitReader(in, 1000, 1000, 100*time.Millisecond)

	expected := []byte("hello world")

	resultCh := make(chan struct{})
	go func() {
		outBytes, err := ioutil.ReadAll(limit)
		if err != nil {
			t.Fatalf("ReadAll failed: %v", err)
		}

		if reflect.DeepEqual(outBytes, expected) {
			close(resultCh)
			return
		}
	}()

	// Send the data
	in.data <- expected
	in.Close()

	select {
	case <-resultCh:
	case <-time.After(1 * time.Second):
		t.Fatalf("did not exit by time limit")
	}
}

const (
	job = `job "job1" {
        type = "service"
        datacenters = [ "dc1" ]
        group "group1" {
                count = 1
                task "task1" {
                        driver = "exec"
                        resources = {}
                }
                restart{
                        attempts = 10
                        mode = "delay"
						interval = "15s"
                }
        }
}`
)

var (
	expectedApiJob = &api.Job{
		ID:          helper.StringToPtr("job1"),
		Name:        helper.StringToPtr("job1"),
		Type:        helper.StringToPtr("service"),
		Datacenters: []string{"dc1"},
		TaskGroups: []*api.TaskGroup{
			{
				Name:  helper.StringToPtr("group1"),
				Count: helper.IntToPtr(1),
				RestartPolicy: &api.RestartPolicy{
					Attempts: helper.IntToPtr(10),
					Interval: helper.TimeToPtr(15 * time.Second),
					Mode:     helper.StringToPtr("delay"),
				},

				Tasks: []*api.Task{
					{
						Driver:    "exec",
						Name:      "task1",
						Resources: &api.Resources{},
					},
				},
			},
		},
	}
)

// Test APIJob with local jobfile
func TestJobGetter_LocalFile(t *testing.T) {
	t.Parallel()
	fh, err := ioutil.TempFile("", "nomad")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.Remove(fh.Name())
	_, err = fh.WriteString(job)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	j := &JobGetter{}
	aj, err := j.ApiJob(fh.Name())
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if !reflect.DeepEqual(expectedApiJob, aj) {
		eflat := flatmap.Flatten(expectedApiJob, nil, false)
		aflat := flatmap.Flatten(aj, nil, false)
		t.Fatalf("got:\n%v\nwant:\n%v", aflat, eflat)
	}
}

// Test StructJob with jobfile from HTTP Server
func TestJobGetter_HTTPServer(t *testing.T) {
	t.Parallel()
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, job)
	})
	go http.ListenAndServe("127.0.0.1:12345", nil)

	// Wait until HTTP Server starts certainly
	time.Sleep(100 * time.Millisecond)

	j := &JobGetter{}
	aj, err := j.ApiJob("http://127.0.0.1:12345/")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if !reflect.DeepEqual(expectedApiJob, aj) {
		for _, d := range pretty.Diff(expectedApiJob, aj) {
			t.Logf(d)
		}
		t.Fatalf("Unexpected file")
	}
}

func TestPrettyTimeDiff(t *testing.T) {
	now := time.Now()
	test_cases := []struct {
		t1  time.Time
		t2  time.Time
		exp string
	}{
		{now, time.Unix(0, 0), ""}, // This is the upgrade path case
		{now, now.Add(-10 * time.Millisecond), "0s ago"},
		{now, now.Add(-10 * time.Millisecond), "0s ago"},
		{now, now.Add(-740 * time.Second), "12m20s ago"},
		{now, now.Add(-12 * time.Minute), "12m ago"},
		{now, now.Add(-60 * time.Minute), "1h ago"},
		{now, now.Add(-80 * time.Minute), "1h20m ago"},
		{now, now.Add(-6 * time.Hour), "6h ago"},
		{now.Add(-6 * time.Hour), now, "6h from now"},
		{now, now.Add(-22165 * time.Second), "6h9m ago"},
		{now, now.Add(-100 * time.Hour), "4d4h ago"},
		{now, now.Add(-438000 * time.Minute), "10mo4d ago"},
		{now, now.Add(-20460 * time.Hour), "2y4mo ago"},
	}
	for _, tc := range test_cases {
		out := prettyTimeDiff(tc.t2, tc.t1)
		if out != tc.exp {
			t.Fatalf("expected :%v but got :%v", tc.exp, out)
		}
	}

	var t1 time.Time
	out := prettyTimeDiff(t1, time.Now())

	if out != "" {
		t.Fatalf("Expected empty output but got:%v", out)
	}

}
<commit_msg>fix flaky pretty time test<commit_after>package command

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/helper"
	"github.com/hashicorp/nomad/helper/flatmap"
	"github.com/kr/pretty"
	"github.com/mitchellh/cli"
)

func TestHelpers_FormatKV(t *testing.T) {
	t.Parallel()
	in := []string{"alpha|beta", "charlie|delta", "echo|"}
	out := formatKV(in)

	expect := "alpha   = beta\n"
	expect += "charlie = delta\n"
	expect += "echo    = <none>"

	if out != expect {
		t.Fatalf("expect: %s, got: %s", expect, out)
	}
}

func TestHelpers_FormatList(t *testing.T) {
	t.Parallel()
	in := []string{"alpha|beta||delta"}
	out := formatList(in)

	expect := "alpha  beta  <none>  delta"

	if out != expect {
		t.Fatalf("expect: %s, got: %s", expect, out)
	}
}

func TestHelpers_NodeID(t *testing.T) {
	t.Parallel()
	srv, _, _ := testServer(t, false, nil)
	defer srv.Shutdown()

	meta := Meta{Ui: new(cli.MockUi)}
	client, err := meta.Client()
	if err != nil {
		t.FailNow()
	}

	// This is because there is no client
	if _, err := getLocalNodeID(client); err == nil {
		t.Fatalf("getLocalNodeID() should fail")
	}
}

func TestHelpers_LineLimitReader_NoTimeLimit(t *testing.T) {
	t.Parallel()
	helloString := `hello
world
this
is
a
test`

	noLines := "jskdfhjasdhfjkajkldsfdlsjkahfkjdsafa"

	cases := []struct {
		Input       string
		Output      string
		Lines       int
		SearchLimit int
	}{
		{
			Input:       helloString,
			Output:      helloString,
			Lines:       6,
			SearchLimit: 1000,
		},
		{
			Input: helloString,
			Output: `world
this
is
a
test`,
			Lines:       5,
			SearchLimit: 1000,
		},
		{
			Input:       helloString,
			Output:      `test`,
			Lines:       1,
			SearchLimit: 1000,
		},
		{
			Input:       helloString,
			Output:      "",
			Lines:       0,
			SearchLimit: 1000,
		},
		{
			Input:       helloString,
			Output:      helloString,
			Lines:       6,
			SearchLimit: 1, // Exceed the limit
		},
		{
			Input:       noLines,
			Output:      noLines,
			Lines:       10,
			SearchLimit: 1000,
		},
		{
			Input:       noLines,
			Output:      noLines,
			Lines:       10,
			SearchLimit: 2,
		},
	}

	for i, c := range cases {
		in := ioutil.NopCloser(strings.NewReader(c.Input))
		limit := NewLineLimitReader(in, c.Lines, c.SearchLimit, 0)
		outBytes, err := ioutil.ReadAll(limit)
		if err != nil {
			t.Fatalf("case %d failed: %v", i, err)
		}

		out := string(outBytes)
		if out != c.Output {
			t.Fatalf("case %d: got %q; want %q", i, out, c.Output)
		}
	}
}

type testReadCloser struct {
	data chan []byte
}

func (t *testReadCloser) Read(p []byte) (n int, err error) {
	select {
	case b, ok := <-t.data:
		if !ok {
			return 0, io.EOF
		}

		return copy(p, b), nil
	case <-time.After(10 * time.Millisecond):
		return 0, nil
	}
}

func (t *testReadCloser) Close() error {
	close(t.data)
	return nil
}

func TestHelpers_LineLimitReader_TimeLimit(t *testing.T) {
	t.Parallel()
	// Create the test reader
	in := &testReadCloser{data: make(chan []byte)}

	// Set up the reader such that it won't hit the line/buffer limit and could
	// only terminate if it hits the time limit
	limit := NewLineLimitReader(in, 1000, 1000, 100*time.Millisecond)

	expected := []byte("hello world")

	resultCh := make(chan struct{})
	go func() {
		outBytes, err := ioutil.ReadAll(limit)
		if err != nil {
			t.Fatalf("ReadAll failed: %v", err)
		}

		if reflect.DeepEqual(outBytes, expected) {
			close(resultCh)
			return
		}
	}()

	// Send the data
	in.data <- expected
	in.Close()

	select {
	case <-resultCh:
	case <-time.After(1 * time.Second):
		t.Fatalf("did not exit by time limit")
	}
}

const (
	job = `job "job1" {
        type = "service"
        datacenters = [ "dc1" ]
        group "group1" {
                count = 1
                task "task1" {
                        driver = "exec"
                        resources = {}
                }
                restart{
                        attempts = 10
                        mode = "delay"
						interval = "15s"
                }
        }
}`
)

var (
	expectedApiJob = &api.Job{
		ID:          helper.StringToPtr("job1"),
		Name:        helper.StringToPtr("job1"),
		Type:        helper.StringToPtr("service"),
		Datacenters: []string{"dc1"},
		TaskGroups: []*api.TaskGroup{
			{
				Name:  helper.StringToPtr("group1"),
				Count: helper.IntToPtr(1),
				RestartPolicy: &api.RestartPolicy{
					Attempts: helper.IntToPtr(10),
					Interval: helper.TimeToPtr(15 * time.Second),
					Mode:     helper.StringToPtr("delay"),
				},

				Tasks: []*api.Task{
					{
						Driver:    "exec",
						Name:      "task1",
						Resources: &api.Resources{},
					},
				},
			},
		},
	}
)

// Test APIJob with local jobfile
func TestJobGetter_LocalFile(t *testing.T) {
	t.Parallel()
	fh, err := ioutil.TempFile("", "nomad")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.Remove(fh.Name())
	_, err = fh.WriteString(job)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	j := &JobGetter{}
	aj, err := j.ApiJob(fh.Name())
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if !reflect.DeepEqual(expectedApiJob, aj) {
		eflat := flatmap.Flatten(expectedApiJob, nil, false)
		aflat := flatmap.Flatten(aj, nil, false)
		t.Fatalf("got:\n%v\nwant:\n%v", aflat, eflat)
	}
}

// Test StructJob with jobfile from HTTP Server
func TestJobGetter_HTTPServer(t *testing.T) {
	t.Parallel()
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, job)
	})
	go http.ListenAndServe("127.0.0.1:12345", nil)

	// Wait until HTTP Server starts certainly
	time.Sleep(100 * time.Millisecond)

	j := &JobGetter{}
	aj, err := j.ApiJob("http://127.0.0.1:12345/")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if !reflect.DeepEqual(expectedApiJob, aj) {
		for _, d := range pretty.Diff(expectedApiJob, aj) {
			t.Logf(d)
		}
		t.Fatalf("Unexpected file")
	}
}

func TestPrettyTimeDiff(t *testing.T) {
	// Grab the time and truncate to the nearest second. This allows our tests
	// to be deterministic since we don't have to worry about rounding.
	now := time.Now().Truncate(time.Second)

	test_cases := []struct {
		t1  time.Time
		t2  time.Time
		exp string
	}{
		{now, time.Unix(0, 0), ""}, // This is the upgrade path case
		{now, now.Add(-10 * time.Millisecond), "0s ago"},
		{now, now.Add(-740 * time.Second), "12m20s ago"},
		{now, now.Add(-12 * time.Minute), "12m ago"},
		{now, now.Add(-60 * time.Minute), "1h ago"},
		{now, now.Add(-80 * time.Minute), "1h20m ago"},
		{now, now.Add(-6 * time.Hour), "6h ago"},
		{now.Add(-6 * time.Hour), now, "6h from now"},
		{now, now.Add(-22165 * time.Second), "6h9m ago"},
		{now, now.Add(-100 * time.Hour), "4d4h ago"},
		{now, now.Add(-438000 * time.Minute), "10mo4d ago"},
		{now, now.Add(-20460 * time.Hour), "2y4mo ago"},
	}
	for _, tc := range test_cases {
		t.Run(tc.exp, func(t *testing.T) {
			out := prettyTimeDiff(tc.t2, tc.t1)
			if out != tc.exp {
				t.Fatalf("expected :%v but got :%v", tc.exp, out)
			}
		})
	}

	var t1 time.Time
	out := prettyTimeDiff(t1, time.Now())

	if out != "" {
		t.Fatalf("Expected empty output but got:%v", out)
	}

}
<|endoftext|>


<commit_before>/*
   Copyright 2013 Juliano Martinez <juliano@martinez.io>

   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.

      Based on http://github.com/nf/webfront

   @author: Juliano Martinez
*/

package http_server

import (
	"errors"
	"fmt"
	"github.com/fiorix/go-redis/redis"
	"github.com/ncode/hot-potato-router/utils"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	cfg = utils.NewConfig()
	rc  = redis.New(cfg.Options["redis"]["server_list"])
)

func xff(req *http.Request) string {
	remote_addr := strings.Split(req.RemoteAddr, ":")
	if len(remote_addr) == 0 {
		return ""
	}
	return remote_addr[0]
}

type Server struct {
	mu      sync.RWMutex
	proxy   map[string][]Proxy
	backend map[string]int
	vcount  map[string]map[string]int
}

type Proxy struct {
	//	last    time.Time
	Backend string
	handler http.Handler
}

func Listen(addr string) net.Listener {
	var l net.Listener
	var err error

	l, err = net.Listen("tcp", addr)
	if err != nil {
		log.Fatal(err)
	}
	return l
}

func NewServer(probe time.Duration) (*Server, error) {
	s := new(Server)
	s.proxy = make(map[string][]Proxy)
	s.backend = make(map[string]int)
	s.vcount = make(map[string]map[string]int)
	go s.probe_backends(probe)
	return s, nil
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if h := s.handler(r); h != nil {
		client := xff(r)
		utils.Log(fmt.Sprintf("Request from: %s Url: %s", client, r.Host))
		r.Header.Add("X-Forwarded-For", client)
		r.Header.Add("X-Real-IP", client)
		h.ServeHTTP(w, r)
		return
	}
	http.Error(w, "Not found.", http.StatusNotFound)
}

func (s *Server) handler(req *http.Request) http.Handler {
	vhost := req.Host
	if i := strings.Index(vhost, ":"); i >= 0 {
		vhost = vhost[:i]
	}

	_, ok := s.proxy[vhost]
	if !ok {
		err := s.populate_proxies(vhost)
		if err != nil {
			utils.Log(fmt.Sprintf("%s for vhost %s", err, vhost))
			return nil
		}
	}
	return s.Next(vhost)
}

func (s *Server) populate_proxies(vhost string) (err error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	f, _ := rc.ZRange(fmt.Sprintf("hpr-backends::%s", vhost), 0, -1, true)
	if len(f) == 0 {
		if len(s.proxy[vhost]) > 0 {
			delete(s.proxy, vhost)
			delete(s.backend, vhost)
			delete(s.vcount, vhost)
		}
		return errors.New("Backend list is empty")
	}

	var url string
	for _, be := range f {
		count, err := strconv.Atoi(be)
		if err != nil {
			url = be
			continue
		}
		backend := fmt.Sprintf("http://%s", url)
		for r := 1; r <= count; r++ {
			if r <= s.vcount[vhost][backend] {
				continue
			}
			s.proxy[vhost] = append(s.proxy[vhost], Proxy{backend, makeHandler(url)})
			v, ready := s.vcount[vhost]
			if !ready {
				v = make(map[string]int)
				s.vcount[vhost] = v
			}
			s.vcount[vhost][backend]++
			utils.Log(fmt.Sprintf("Backend %s with %d handlers for %s", backend, s.vcount[vhost][backend], vhost))
		}
	}
	return
}

/* TODO: Implement more balance algorithms */
func (s *Server) Next(vhost string) http.Handler {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.backend[vhost]++
	total := len(s.proxy[vhost])
	if s.backend[vhost] >= total {
		s.backend[vhost] = 0
	}
	utils.Log(fmt.Sprintf(
		"Using backend: %s Url: %s", s.proxy[vhost][s.backend[vhost]].Backend, vhost))
	return s.proxy[vhost][s.backend[vhost]].handler
}

/* TODO: Implement more probes */
func (s *Server) probe_backends(probe time.Duration) {
	transport := http.Transport{Dial: dialTimeout}
	client := &http.Client{
		Transport: &transport,
	}

	for {
		time.Sleep(probe)
		for vhost, backends := range s.proxy {
			fmt.Println(backends)
			err := s.populate_proxies(vhost)
			if err != nil {
				utils.Log(fmt.Sprintf("Cleaned entries from vhost: %s", vhost))
				continue
			}
			is_dead := make(map[string]bool)
			removed := 0
			for backend := range backends {
				backend = backend - removed
				utils.Log(fmt.Sprintf(
					"vhost: %s backends: %s", vhost, s.proxy[vhost][backend].Backend))
				if is_dead[s.proxy[vhost][backend].Backend] == true {
					utils.Log(fmt.Sprintf("Removing dead backend: %s", s.proxy[vhost][backend].Backend))
					s.mu.Lock()
					s.proxy[vhost] = s.proxy[vhost][:backend+copy(s.proxy[vhost][backend:], s.proxy[vhost][backend+1:])]
					s.vcount[vhost][s.proxy[vhost][backend].Backend]--
					s.mu.Unlock()
					removed++
					continue
				}

				req, err := client.Get(s.proxy[vhost][backend].Backend)
				defer resp.Body.Close()
				if err != nil {
					utils.Log(fmt.Sprintf("Removing dead backend: %s with error %s", s.proxy[vhost][backend].Backend, err))
					is_dead[s.proxy[vhost][backend].Backend] = true
					s.mu.Lock()
					s.proxy[vhost] = s.proxy[vhost][:backend+copy(s.proxy[vhost][backend:], s.proxy[vhost][backend+1:])]
					s.vcount[vhost][s.proxy[vhost][backend].Backend]--
					s.mu.Unlock()
					removed++
				} else {
					utils.Log(fmt.Sprintf("Alive: %s", s.proxy[vhost][backend].Backend))
				}
			}

		}

	}
}

func dialTimeout(network, addr string) (net.Conn, error) {
	timeout := time.Duration(2 * time.Second)
	return net.DialTimeout(network, addr, timeout)
}

func makeHandler(f string) http.Handler {
	if f != "" {
		return &httputil.ReverseProxy{
			Director: func(req *http.Request) {
				req.URL.Scheme = "http"
				req.URL.Host = f
			},
		}
	}
	return nil
}
<commit_msg>killing probes after used<commit_after>/*
   Copyright 2013 Juliano Martinez <juliano@martinez.io>

   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.

      Based on http://github.com/nf/webfront

   @author: Juliano Martinez
*/

package http_server

import (
	"errors"
	"fmt"
	"github.com/fiorix/go-redis/redis"
	"github.com/ncode/hot-potato-router/utils"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	cfg = utils.NewConfig()
	rc  = redis.New(cfg.Options["redis"]["server_list"])
)

func xff(req *http.Request) string {
	remote_addr := strings.Split(req.RemoteAddr, ":")
	if len(remote_addr) == 0 {
		return ""
	}
	return remote_addr[0]
}

type Server struct {
	mu      sync.RWMutex
	proxy   map[string][]Proxy
	backend map[string]int
	vcount  map[string]map[string]int
}

type Proxy struct {
	//	last    time.Time
	Backend string
	handler http.Handler
}

func Listen(addr string) net.Listener {
	var l net.Listener
	var err error

	l, err = net.Listen("tcp", addr)
	if err != nil {
		log.Fatal(err)
	}
	return l
}

func NewServer(probe time.Duration) (*Server, error) {
	s := new(Server)
	s.proxy = make(map[string][]Proxy)
	s.backend = make(map[string]int)
	s.vcount = make(map[string]map[string]int)
	go s.probe_backends(probe)
	return s, nil
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if h := s.handler(r); h != nil {
		client := xff(r)
		utils.Log(fmt.Sprintf("Request from: %s Url: %s", client, r.Host))
		r.Header.Add("X-Forwarded-For", client)
		r.Header.Add("X-Real-IP", client)
		h.ServeHTTP(w, r)
		return
	}
	http.Error(w, "Not found.", http.StatusNotFound)
}

func (s *Server) handler(req *http.Request) http.Handler {
	vhost := req.Host
	if i := strings.Index(vhost, ":"); i >= 0 {
		vhost = vhost[:i]
	}

	_, ok := s.proxy[vhost]
	if !ok {
		err := s.populate_proxies(vhost)
		if err != nil {
			utils.Log(fmt.Sprintf("%s for vhost %s", err, vhost))
			return nil
		}
	}
	return s.Next(vhost)
}

func (s *Server) populate_proxies(vhost string) (err error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	f, _ := rc.ZRange(fmt.Sprintf("hpr-backends::%s", vhost), 0, -1, true)
	if len(f) == 0 {
		if len(s.proxy[vhost]) > 0 {
			delete(s.proxy, vhost)
			delete(s.backend, vhost)
			delete(s.vcount, vhost)
		}
		return errors.New("Backend list is empty")
	}

	var url string
	for _, be := range f {
		count, err := strconv.Atoi(be)
		if err != nil {
			url = be
			continue
		}
		backend := fmt.Sprintf("http://%s", url)
		for r := 1; r <= count; r++ {
			if r <= s.vcount[vhost][backend] {
				continue
			}
			s.proxy[vhost] = append(s.proxy[vhost], Proxy{backend, makeHandler(url)})
			v, ready := s.vcount[vhost]
			if !ready {
				v = make(map[string]int)
				s.vcount[vhost] = v
			}
			s.vcount[vhost][backend]++
			utils.Log(fmt.Sprintf("Backend %s with %d handlers for %s", backend, s.vcount[vhost][backend], vhost))
		}
	}
	return
}

/* TODO: Implement more balance algorithms */
func (s *Server) Next(vhost string) http.Handler {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.backend[vhost]++
	total := len(s.proxy[vhost])
	if s.backend[vhost] >= total {
		s.backend[vhost] = 0
	}
	utils.Log(fmt.Sprintf(
		"Using backend: %s Url: %s", s.proxy[vhost][s.backend[vhost]].Backend, vhost))
	return s.proxy[vhost][s.backend[vhost]].handler
}

/* TODO: Implement more probes */
func (s *Server) probe_backends(probe time.Duration) {
	transport := http.Transport{Dial: dialTimeout}
	client := &http.Client{
		Transport: &transport,
	}

	for {
		time.Sleep(probe)
		for vhost, backends := range s.proxy {
			fmt.Println(backends)
			err := s.populate_proxies(vhost)
			if err != nil {
				utils.Log(fmt.Sprintf("Cleaned entries from vhost: %s", vhost))
				continue
			}
			is_dead := make(map[string]bool)
			removed := 0
			for backend := range backends {
				backend = backend - removed
				utils.Log(fmt.Sprintf(
					"vhost: %s backends: %s", vhost, s.proxy[vhost][backend].Backend))
				if is_dead[s.proxy[vhost][backend].Backend] == true {
					utils.Log(fmt.Sprintf("Removing dead backend: %s", s.proxy[vhost][backend].Backend))
					s.mu.Lock()
					s.proxy[vhost] = s.proxy[vhost][:backend+copy(s.proxy[vhost][backend:], s.proxy[vhost][backend+1:])]
					s.vcount[vhost][s.proxy[vhost][backend].Backend]--
					s.mu.Unlock()
					removed++
					continue
				}

				req, err := client.Get(s.proxy[vhost][backend].Backend)
				if err != nil {
					utils.Log(fmt.Sprintf("Removing dead backend: %s with error %s", s.proxy[vhost][backend].Backend, err))
					is_dead[s.proxy[vhost][backend].Backend] = true
					s.mu.Lock()
					s.proxy[vhost] = s.proxy[vhost][:backend+copy(s.proxy[vhost][backend:], s.proxy[vhost][backend+1:])]
					s.vcount[vhost][s.proxy[vhost][backend].Backend]--
					s.mu.Unlock()
					removed++
				} else {
					utils.Log(fmt.Sprintf("Alive: %s", s.proxy[vhost][backend].Backend))
				}
				transport.CancelRequest(req)
			}

		}

	}
}

func dialTimeout(network, addr string) (net.Conn, error) {
	timeout := time.Duration(2 * time.Second)
	return net.DialTimeout(network, addr, timeout)
}

func makeHandler(f string) http.Handler {
	if f != "" {
		return &httputil.ReverseProxy{
			Director: func(req *http.Request) {
				req.URL.Scheme = "http"
				req.URL.Host = f
			},
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package http

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"reflect"
	"strconv"
	"strings"

	cmds "github.com/ipfs/go-ipfs/commands"
	config "github.com/ipfs/go-ipfs/repo/config"

	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
)

const (
	ApiUrlFormat = "http://%s%s/%s?%s"
	ApiPath      = "/api/v0" // TODO: make configurable
)

// Client is the commands HTTP client interface.
type Client interface {
	Send(req cmds.Request) (cmds.Response, error)
}

type client struct {
	serverAddress string
}

func NewClient(address string) Client {
	return &client{address}
}

func (c *client) Send(req cmds.Request) (cmds.Response, error) {

	if req.Context() == nil {
		log.Warningf("no context set in request")
		err := req.SetRootContext(context.TODO())
		if err != nil {
			return nil, err
		}
	}

	// save user-provided encoding
	previousUserProvidedEncoding, found, err := req.Option(cmds.EncShort).String()
	if err != nil {
		return nil, err
	}

	// override with json to send to server
	req.SetOption(cmds.EncShort, cmds.JSON)

	// stream channel output
	req.SetOption(cmds.ChanOpt, "true")

	query, err := getQuery(req)
	if err != nil {
		return nil, err
	}

	var fileReader *MultiFileReader
	var reader io.Reader

	if req.Files() != nil {
		fileReader = NewMultiFileReader(req.Files(), true)
		reader = fileReader
	} else {
		// if we have no file data, use an empty Reader
		// (http.NewRequest panics when a nil Reader is used)
		reader = strings.NewReader("")
	}

	path := strings.Join(req.Path(), "/")
	url := fmt.Sprintf(ApiUrlFormat, c.serverAddress, ApiPath, path, query)

	httpReq, err := http.NewRequest("POST", url, reader)
	if err != nil {
		return nil, err
	}

	// TODO extract string consts?
	if fileReader != nil {
		httpReq.Header.Set("Content-Type", "multipart/form-data; boundary="+fileReader.Boundary())
		httpReq.Header.Set("Content-Disposition", "form-data: name=\"files\"")
	} else {
		httpReq.Header.Set("Content-Type", "application/octet-stream")
	}
	version := config.CurrentVersionNumber
	httpReq.Header.Set("User-Agent", fmt.Sprintf("/go-ipfs/%s/", version))

	ec := make(chan error, 1)
	rc := make(chan cmds.Response, 1)
	dc := req.Context().Done()

	go func() {
		httpRes, err := http.DefaultClient.Do(httpReq)
		if err != nil {
			ec <- err
			return
		}

		// using the overridden JSON encoding in request
		res, err := getResponse(httpRes, req)
		if err != nil {
			ec <- err
			return
		}

		rc <- res
	}()

	for {
		select {
		case <-dc:
			log.Debug("Context cancelled, cancelling HTTP request...")
			tr := http.DefaultTransport.(*http.Transport)
			tr.CancelRequest(httpReq)
			dc = nil // Wait for ec or rc
		case err := <-ec:
			return nil, err
		case res := <-rc:
			if found && len(previousUserProvidedEncoding) > 0 {
				// reset to user provided encoding after sending request
				// NB: if user has provided an encoding but it is the empty string,
				// still leave it as JSON.
				req.SetOption(cmds.EncShort, previousUserProvidedEncoding)
			}
			return res, nil
		}
	}
}

func getQuery(req cmds.Request) (string, error) {
	query := url.Values{}
	for k, v := range req.Options() {
		str := fmt.Sprintf("%v", v)
		query.Set(k, str)
	}

	args := req.Arguments()
	argDefs := req.Command().Arguments

	argDefIndex := 0

	for _, arg := range args {
		argDef := argDefs[argDefIndex]
		// skip ArgFiles
		for argDef.Type == cmds.ArgFile {
			argDefIndex++
			argDef = argDefs[argDefIndex]
		}

		query.Add("arg", arg)

		if len(argDefs) > argDefIndex+1 {
			argDefIndex++
		}
	}

	return query.Encode(), nil
}

// getResponse decodes a http.Response to create a cmds.Response
func getResponse(httpRes *http.Response, req cmds.Request) (cmds.Response, error) {
	var err error
	res := cmds.NewResponse(req)

	contentType := httpRes.Header.Get(contentTypeHeader)
	contentType = strings.Split(contentType, ";")[0]

	lengthHeader := httpRes.Header.Get(contentLengthHeader)
	if len(lengthHeader) > 0 {
		length, err := strconv.ParseUint(lengthHeader, 10, 64)
		if err != nil {
			return nil, err
		}
		res.SetLength(length)
	}

	res.SetCloser(httpRes.Body)

	if len(httpRes.Header.Get(streamHeader)) > 0 && contentType != "application/json" {
		// if output is a stream, we can just use the body reader
		res.SetOutput(&httpResponseReader{httpRes})
		return res, nil

	} else if len(httpRes.Header.Get(channelHeader)) > 0 {
		// if output is coming from a channel, decode each chunk
		outChan := make(chan interface{})
		go func() {
			dec := json.NewDecoder(&httpResponseReader{httpRes})
			outputType := reflect.TypeOf(req.Command().Type)

			ctx := req.Context()

			for {
				var v interface{}
				var err error
				if outputType != nil {
					v = reflect.New(outputType).Interface()
					err = dec.Decode(v)
				} else {
					err = dec.Decode(&v)
				}

				// since we are just looping reading on the response, the only way to
				// know we are 'done' is for the consumer to close the response body.
				// doing so doesnt throw an io.EOF, but we want to treat it like one.
				if err != nil && strings.Contains(err.Error(), "read on closed response body") {
					err = io.EOF
				}
				if err != nil && err != io.EOF {
					log.Error(err)
					return
				}

				select {
				case <-ctx.Done():
					close(outChan)
					return
				default:
				}

				if err == io.EOF {
					close(outChan)
					return
				}
				outChan <- v
			}
		}()

		res.SetOutput((<-chan interface{})(outChan))
		return res, nil
	}

	dec := json.NewDecoder(&httpResponseReader{httpRes})

	if httpRes.StatusCode >= http.StatusBadRequest {
		e := cmds.Error{}

		if httpRes.StatusCode == http.StatusNotFound {
			// handle 404s
			e.Message = "Command not found."
			e.Code = cmds.ErrClient

		} else if contentType == "text/plain" {
			// handle non-marshalled errors
			buf := bytes.NewBuffer(nil)
			io.Copy(buf, httpRes.Body)
			e.Message = string(buf.Bytes())
			e.Code = cmds.ErrNormal

		} else {
			// handle marshalled errors
			err = dec.Decode(&e)
			if err != nil {
				return nil, err
			}
		}

		res.SetError(e, e.Code)

	} else {
		outputType := reflect.TypeOf(req.Command().Type)
		var v interface{}

		if outputType != nil {
			v = reflect.New(outputType).Interface()
			err = dec.Decode(v)
		} else {
			err = dec.Decode(&v)
		}
		if err != nil && err != io.EOF {
			return nil, err
		}
		if v != nil {
			res.SetOutput(v)
		}
	}

	return res, nil
}

type httpResponseReader struct {
	resp *http.Response
}

func (r *httpResponseReader) Read(b []byte) (int, error) {
	n, err := r.resp.Body.Read(b)
	if err == io.EOF {
		_ = r.resp.Body.Close()
		trailerErr := r.checkError()
		if trailerErr != nil {
			return n, trailerErr
		}
	}
	return n, err
}

func (r *httpResponseReader) checkError() error {
	if e := r.resp.Trailer.Get(StreamErrHeader); e != "" {
		return errors.New(e)
	}
	return nil
}

func (r *httpResponseReader) Close() error {
	return r.resp.Body.Close()
}
<commit_msg>refactor http client code<commit_after>package http

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"reflect"
	"strconv"
	"strings"

	cmds "github.com/ipfs/go-ipfs/commands"
	config "github.com/ipfs/go-ipfs/repo/config"

	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
)

const (
	ApiUrlFormat = "http://%s%s/%s?%s"
	ApiPath      = "/api/v0" // TODO: make configurable
)

// Client is the commands HTTP client interface.
type Client interface {
	Send(req cmds.Request) (cmds.Response, error)
}

type client struct {
	serverAddress string
}

func NewClient(address string) Client {
	return &client{address}
}

func (c *client) Send(req cmds.Request) (cmds.Response, error) {

	if req.Context() == nil {
		log.Warningf("no context set in request")
		err := req.SetRootContext(context.TODO())
		if err != nil {
			return nil, err
		}
	}

	// save user-provided encoding
	previousUserProvidedEncoding, found, err := req.Option(cmds.EncShort).String()
	if err != nil {
		return nil, err
	}

	// override with json to send to server
	req.SetOption(cmds.EncShort, cmds.JSON)

	// stream channel output
	req.SetOption(cmds.ChanOpt, "true")

	query, err := getQuery(req)
	if err != nil {
		return nil, err
	}

	var fileReader *MultiFileReader
	var reader io.Reader

	if req.Files() != nil {
		fileReader = NewMultiFileReader(req.Files(), true)
		reader = fileReader
	} else {
		// if we have no file data, use an empty Reader
		// (http.NewRequest panics when a nil Reader is used)
		reader = strings.NewReader("")
	}

	path := strings.Join(req.Path(), "/")
	url := fmt.Sprintf(ApiUrlFormat, c.serverAddress, ApiPath, path, query)

	httpReq, err := http.NewRequest("POST", url, reader)
	if err != nil {
		return nil, err
	}

	// TODO extract string consts?
	if fileReader != nil {
		httpReq.Header.Set("Content-Type", "multipart/form-data; boundary="+fileReader.Boundary())
		httpReq.Header.Set("Content-Disposition", "form-data: name=\"files\"")
	} else {
		httpReq.Header.Set("Content-Type", "application/octet-stream")
	}
	version := config.CurrentVersionNumber
	httpReq.Header.Set("User-Agent", fmt.Sprintf("/go-ipfs/%s/", version))

	ec := make(chan error, 1)
	rc := make(chan cmds.Response, 1)
	dc := req.Context().Done()

	go func() {
		httpRes, err := http.DefaultClient.Do(httpReq)
		if err != nil {
			ec <- err
			return
		}

		// using the overridden JSON encoding in request
		res, err := getResponse(httpRes, req)
		if err != nil {
			ec <- err
			return
		}

		rc <- res
	}()

	for {
		select {
		case <-dc:
			log.Debug("Context cancelled, cancelling HTTP request...")
			tr := http.DefaultTransport.(*http.Transport)
			tr.CancelRequest(httpReq)
			dc = nil // Wait for ec or rc
		case err := <-ec:
			return nil, err
		case res := <-rc:
			if found && len(previousUserProvidedEncoding) > 0 {
				// reset to user provided encoding after sending request
				// NB: if user has provided an encoding but it is the empty string,
				// still leave it as JSON.
				req.SetOption(cmds.EncShort, previousUserProvidedEncoding)
			}
			return res, nil
		}
	}
}

func getQuery(req cmds.Request) (string, error) {
	query := url.Values{}
	for k, v := range req.Options() {
		str := fmt.Sprintf("%v", v)
		query.Set(k, str)
	}

	args := req.Arguments()
	argDefs := req.Command().Arguments

	argDefIndex := 0

	for _, arg := range args {
		argDef := argDefs[argDefIndex]
		// skip ArgFiles
		for argDef.Type == cmds.ArgFile {
			argDefIndex++
			argDef = argDefs[argDefIndex]
		}

		query.Add("arg", arg)

		if len(argDefs) > argDefIndex+1 {
			argDefIndex++
		}
	}

	return query.Encode(), nil
}

// getResponse decodes a http.Response to create a cmds.Response
func getResponse(httpRes *http.Response, req cmds.Request) (cmds.Response, error) {
	var err error
	res := cmds.NewResponse(req)

	contentType := httpRes.Header.Get(contentTypeHeader)
	contentType = strings.Split(contentType, ";")[0]

	lengthHeader := httpRes.Header.Get(contentLengthHeader)
	if len(lengthHeader) > 0 {
		length, err := strconv.ParseUint(lengthHeader, 10, 64)
		if err != nil {
			return nil, err
		}
		res.SetLength(length)
	}

	res.SetCloser(httpRes.Body)

	if contentType != "application/json" {
		// for all non json output types, just stream back the output
		res.SetOutput(&httpResponseReader{httpRes})
		return res, nil

	} else if len(httpRes.Header.Get(channelHeader)) > 0 {
		// if output is coming from a channel, decode each chunk
		outChan := make(chan interface{})

		go readStreamedJson(req, httpRes, outChan)

		res.SetOutput((<-chan interface{})(outChan))
		return res, nil
	}

	dec := json.NewDecoder(&httpResponseReader{httpRes})

	// If we ran into an error
	if httpRes.StatusCode >= http.StatusBadRequest {
		e := cmds.Error{}

		switch {
		case httpRes.StatusCode == http.StatusNotFound:
			// handle 404s
			e.Message = "Command not found."
			e.Code = cmds.ErrClient

		case contentType == "text/plain":
			// handle non-marshalled errors
			buf := bytes.NewBuffer(nil)
			io.Copy(buf, httpRes.Body)
			e.Message = string(buf.Bytes())
			e.Code = cmds.ErrNormal

		default:
			// handle marshalled errors
			err = dec.Decode(&e)
			if err != nil {
				return nil, err
			}
		}

		res.SetError(e, e.Code)

		return res, nil
	}

	outputType := reflect.TypeOf(req.Command().Type)
	v, err := decodeTypedVal(outputType, dec)
	if err != nil && err != io.EOF {
		return nil, err
	}

	res.SetOutput(v)

	return res, nil
}

func readStreamedJson(req cmds.Request, httpRes *http.Response, out chan<- interface{}) {
	defer close(out)
	dec := json.NewDecoder(&httpResponseReader{httpRes})
	outputType := reflect.TypeOf(req.Command().Type)

	ctx := req.Context()

	for {
		v, err := decodeTypedVal(outputType, dec)
		if err != nil {
			// since we are just looping reading on the response, the only way to
			// know we are 'done' is for the consumer to close the response body.
			// doing so doesnt throw an io.EOF, but we want to treat it like one.
			if !(strings.Contains(err.Error(), "read on closed response body") || err == io.EOF) {
				log.Error(err)
			}
			return
		}

		select {
		case <-ctx.Done():
			return
		case out <- v:
		}

	}
}

func decodeTypedVal(t reflect.Type, dec *json.Decoder) (interface{}, error) {
	var v interface{}
	var err error
	if t != nil {
		v = reflect.New(t).Interface()
		err = dec.Decode(v)
	} else {
		err = dec.Decode(&v)
	}

	return v, err
}

type httpResponseReader struct {
	resp *http.Response
}

func (r *httpResponseReader) Read(b []byte) (int, error) {
	n, err := r.resp.Body.Read(b)
	if err == io.EOF {
		_ = r.resp.Body.Close()
		trailerErr := r.checkError()
		if trailerErr != nil {
			return n, trailerErr
		}
	}
	return n, err
}

func (r *httpResponseReader) checkError() error {
	if e := r.resp.Trailer.Get(StreamErrHeader); e != "" {
		return errors.New(e)
	}
	return nil
}

func (r *httpResponseReader) Close() error {
	return r.resp.Body.Close()
}
<|endoftext|>


<commit_before>package jsoniter

import (
	"encoding/json"
	"fmt"
	"testing"
	"unsafe"

	"github.com/stretchr/testify/require"
)

func Test_write_array_of_interface(t *testing.T) {
	should := require.New(t)
	array := []interface{}{"hello"}
	str, err := MarshalToString(array)
	should.Nil(err)
	should.Equal(`["hello"]`, str)
}

func Test_write_map_of_interface(t *testing.T) {
	should := require.New(t)
	val := map[string]interface{}{"hello": "world"}
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`{"hello":"world"}`, str)
}

func Test_write_map_of_interface_in_struct(t *testing.T) {
	type TestObject struct {
		Field map[string]interface{}
	}
	should := require.New(t)
	val := TestObject{map[string]interface{}{"hello": "world"}}
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`{"Field":{"hello":"world"}}`, str)
}

func Test_write_map_of_interface_in_struct_with_two_fields(t *testing.T) {
	type TestObject struct {
		Field  map[string]interface{}
		Field2 string
	}
	should := require.New(t)
	val := TestObject{map[string]interface{}{"hello": "world"}, ""}
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Contains(str, `"Field":{"hello":"world"}`)
}

type MyInterface interface {
	Hello() string
}

type MyString string

func (ms MyString) Hello() string {
	return string(ms)
}

func Test_write_map_of_custom_interface(t *testing.T) {
	should := require.New(t)
	myStr := MyString("world")
	should.Equal("world", myStr.Hello())
	val := map[string]MyInterface{"hello": myStr}
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`{"hello":"world"}`, str)
}

func Test_write_interface(t *testing.T) {
	should := require.New(t)
	var val interface{}
	val = "hello"
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`"hello"`, str)
}

func Test_read_interface(t *testing.T) {
	should := require.New(t)
	var val interface{}
	err := UnmarshalFromString(`"hello"`, &val)
	should.Nil(err)
	should.Equal("hello", val)
	err = UnmarshalFromString(`1e1`, &val)
	should.Nil(err)
	should.Equal(float64(10), val)
	err = UnmarshalFromString(`1.0e1`, &val)
	should.Nil(err)
	should.Equal(float64(10), val)
	err = json.Unmarshal([]byte(`1.0e1`), &val)
	should.Nil(err)
	should.Equal(float64(10), val)
}

func Test_read_custom_interface(t *testing.T) {
	should := require.New(t)
	var val MyInterface
	RegisterTypeDecoderFunc("jsoniter.MyInterface", func(ptr unsafe.Pointer, iter *Iterator) {
		*((*MyInterface)(ptr)) = MyString(iter.ReadString())
	})
	err := UnmarshalFromString(`"hello"`, &val)
	should.Nil(err)
	should.Equal("hello", val.Hello())
}

func Test_decode_object_contain_empty_interface(t *testing.T) {
	type TestObject struct {
		Field interface{}
	}
	should := require.New(t)
	obj := TestObject{}
	obj.Field = 1024
	should.Nil(UnmarshalFromString(`{"Field": "hello"}`, &obj))
	should.Equal("hello", obj.Field)
}

func Test_decode_object_contain_non_empty_interface(t *testing.T) {
	type TestObject struct {
		Field MyInterface
	}
	should := require.New(t)
	obj := TestObject{}
	obj.Field = MyString("abc")
	should.Nil(UnmarshalFromString(`{"Field": "hello"}`, &obj))
	should.Equal(MyString("hello"), obj.Field)
}

func Test_encode_object_contain_empty_interface(t *testing.T) {
	type TestObject struct {
		Field interface{}
	}
	should := require.New(t)
	obj := TestObject{}
	obj.Field = 1024
	str, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"Field":1024}`, str)
}

func Test_encode_object_contain_non_empty_interface(t *testing.T) {
	type TestObject struct {
		Field MyInterface
	}
	should := require.New(t)
	obj := TestObject{}
	obj.Field = MyString("hello")
	str, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"Field":"hello"}`, str)
}

func Test_nil_non_empty_interface(t *testing.T) {
	ConfigDefault.(*frozenConfig).cleanEncoders()
	ConfigDefault.(*frozenConfig).cleanDecoders()
	type TestObject struct {
		Field []MyInterface
	}
	should := require.New(t)
	obj := TestObject{}
	b := []byte(`{"Field":["AAA"]}`)
	should.NotNil(json.Unmarshal(b, &obj))
	should.NotNil(Unmarshal(b, &obj))
}

func Test_read_large_number_as_interface(t *testing.T) {
	should := require.New(t)
	var val interface{}
	err := Config{UseNumber: true}.Froze().UnmarshalFromString(`123456789123456789123456789`, &val)
	should.Nil(err)
	output, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`123456789123456789123456789`, output)
}

func Test_nested_one_field_struct(t *testing.T) {
	should := require.New(t)
	type YetYetAnotherObject struct {
		Field string
	}
	type YetAnotherObject struct {
		Field *YetYetAnotherObject
	}
	type AnotherObject struct {
		Field *YetAnotherObject
	}
	type TestObject struct {
		Me *AnotherObject
	}
	obj := TestObject{&AnotherObject{&YetAnotherObject{&YetYetAnotherObject{"abc"}}}}
	str, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"Me":{"Field":{"Field":{"Field":"abc"}}}}`, str)
	str, err = MarshalToString(&obj)
	should.Nil(err)
	should.Equal(`{"Me":{"Field":{"Field":{"Field":"abc"}}}}`, str)
}

func Test_struct_with_embedded_ptr_with_tag(t *testing.T) {
	type O1 struct {
		O1F string
	}

	type Option struct {
		O1 *O1
	}

	type T struct {
		Option `json:","`
	}
	var obj T
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"O1":null}`, output)
}

func Test_struct_with_one_nil(t *testing.T) {
	type TestObject struct {
		F *float64
	}
	var obj TestObject
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"F":null}`, output)
}

func Test_struct_with_one_nil_embedded(t *testing.T) {
	type Parent struct {
		Field1 string
		Field2 string
	}
	type TestObject struct {
		*Parent
	}
	obj := TestObject{}
	should := require.New(t)
	bytes, err := json.Marshal(obj)
	should.Nil(err)
	should.Equal("{}", string(bytes))
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{}`, output)
}

func Test_struct_with_not_nil_embedded(t *testing.T) {
	type Parent struct {
		Field0 string
		Field1 []string
		Field2 map[string]interface{}
	}
	type TestObject struct {
		*Parent
	}
	should := require.New(t)
	var obj TestObject
	err := UnmarshalFromString(`{"Field0":"1","Field1":null,"Field2":{"K":"V"}}`, &obj)
	should.Nil(err)
	should.Nil(obj.Field1)
	should.Equal(map[string]interface{}{"K": "V"}, obj.Field2)
	should.Equal("1", obj.Field0)
}

func Test_array_with_one_nil_ptr(t *testing.T) {
	obj := [1]*float64{nil}
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`[null]`, output)
}

func Test_array_with_one_not_nil_ptr(t *testing.T) {
	two := float64(2)
	obj := [1]*float64{&two}
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`[2]`, output)
}

func Test_embedded_array_with_one_nil(t *testing.T) {
	type TestObject struct {
		Field1 int
		Field2 [1]*float64
	}
	var obj TestObject
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Contains(output, `"Field2":[null]`)
}

func Test_array_with_nothing(t *testing.T) {
	var obj [2]*float64
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`[null,null]`, output)
}

func Test_unmarshal_ptr_to_interface(t *testing.T) {
	type TestData struct {
		Name string `json:"name"`
	}
	should := require.New(t)
	var obj interface{} = &TestData{}
	err := json.Unmarshal([]byte(`{"name":"value"}`), &obj)
	should.Nil(err)
	should.Equal("&{value}", fmt.Sprintf("%v", obj))
	obj = interface{}(&TestData{})
	err = Unmarshal([]byte(`{"name":"value"}`), &obj)
	should.Nil(err)
	should.Equal("&{value}", fmt.Sprintf("%v", obj))
}

func Test_nil_out_null_interface(t *testing.T) {
	type TestData struct {
		Field interface{} `json:"field"`
	}
	should := require.New(t)

	var boolVar bool
	obj := TestData{
		Field: &boolVar,
	}

	data1 := []byte(`{"field": true}`)

	err := Unmarshal(data1, &obj)
	should.Equal(nil, err)
	should.Equal(true, *(obj.Field.(*bool)))

	data2 := []byte(`{"field": null}`)

	err = Unmarshal(data2, &obj)
	should.Equal(nil, err)
	should.Equal(nil, obj.Field)

	// Checking stdlib behavior matches.
	obj2 := TestData{
		Field: &boolVar,
	}

	err = json.Unmarshal(data1, &obj2)
	should.Equal(nil, err)
	should.Equal(true, *(obj2.Field.(*bool)))

	err = json.Unmarshal(data2, &obj2)
	should.Equal(nil, err)
	should.Equal(nil, obj2.Field)
}

func Test_omitempty_nil_interface(t *testing.T) {
	type TestData struct {
		Field interface{} `json:"field,omitempty"`
	}
	should := require.New(t)

	obj := TestData{
		Field: nil,
	}

	js, err := json.Marshal(obj)
	should.Equal(nil, err)
	should.Equal("{}", string(js))

	str, err := MarshalToString(obj)
	should.Equal(nil, err)
	should.Equal(string(js), str)
}

func Test_omitempty_nil_nonempty_interface(t *testing.T) {
	type TestData struct {
		Field MyInterface `json:"field,omitempty"`
	}
	should := require.New(t)

	obj := TestData{
		Field: nil,
	}

	js, err := json.Marshal(obj)
	should.Equal(nil, err)
	should.Equal("{}", string(js))

	str, err := MarshalToString(obj)
	should.Equal(nil, err)
	should.Equal(string(js), str)

	err = Unmarshal(js, &obj)
	should.Equal(nil, err)
	should.Equal(nil, obj.Field)
}

func Test_marshal_nil_marshaler_interface(t *testing.T) {
	type TestData struct {
		Field json.Marshaler `json:"field"`
	}
	should := require.New(t)

	obj := TestData{
		Field: nil,
	}

	js, err := json.Marshal(obj)
	should.Equal(nil, err)
	should.Equal(`{"field":null}`, string(js))

	str, err := MarshalToString(obj)
	should.Equal(nil, err)
	should.Equal(string(js), str)
}

func Test_marshal_nil_nonempty_interface(t *testing.T) {
	type TestData struct {
		Field MyInterface `json:"field"`
	}
	should := require.New(t)

	obj := TestData{
		Field: nil,
	}

	js, err := json.Marshal(obj)
	should.Equal(nil, err)
	should.Equal(`{"field":null}`, string(js))

	str, err := MarshalToString(obj)
	should.Equal(nil, err)
	should.Equal(string(js), str)

	obj.Field = MyString("hello")
	err = Unmarshal(js, &obj)
	should.Equal(nil, err)
	should.Equal(nil, obj.Field)
}
<commit_msg>improve test<commit_after>package jsoniter

import (
	"encoding/json"
	"fmt"
	"testing"
	"unsafe"

	"github.com/stretchr/testify/require"
)

func Test_write_array_of_interface(t *testing.T) {
	should := require.New(t)
	array := []interface{}{"hello"}
	str, err := MarshalToString(array)
	should.Nil(err)
	should.Equal(`["hello"]`, str)
}

func Test_write_map_of_interface(t *testing.T) {
	should := require.New(t)
	val := map[string]interface{}{"hello": "world"}
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`{"hello":"world"}`, str)
}

func Test_write_map_of_interface_in_struct(t *testing.T) {
	type TestObject struct {
		Field map[string]interface{}
	}
	should := require.New(t)
	val := TestObject{map[string]interface{}{"hello": "world"}}
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`{"Field":{"hello":"world"}}`, str)
}

func Test_write_map_of_interface_in_struct_with_two_fields(t *testing.T) {
	type TestObject struct {
		Field  map[string]interface{}
		Field2 string
	}
	should := require.New(t)
	val := TestObject{map[string]interface{}{"hello": "world"}, ""}
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Contains(str, `"Field":{"hello":"world"}`)
}

type MyInterface interface {
	Hello() string
}

type MyString string

func (ms MyString) Hello() string {
	return string(ms)
}

func Test_write_map_of_custom_interface(t *testing.T) {
	should := require.New(t)
	myStr := MyString("world")
	should.Equal("world", myStr.Hello())
	val := map[string]MyInterface{"hello": myStr}
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`{"hello":"world"}`, str)
}

func Test_write_interface(t *testing.T) {
	should := require.New(t)
	var val interface{}
	val = "hello"
	str, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`"hello"`, str)
}

func Test_read_interface(t *testing.T) {
	should := require.New(t)
	var val interface{}
	err := UnmarshalFromString(`"hello"`, &val)
	should.Nil(err)
	should.Equal("hello", val)
	err = UnmarshalFromString(`1e1`, &val)
	should.Nil(err)
	should.Equal(float64(10), val)
	err = UnmarshalFromString(`1.0e1`, &val)
	should.Nil(err)
	should.Equal(float64(10), val)
	err = json.Unmarshal([]byte(`1.0e1`), &val)
	should.Nil(err)
	should.Equal(float64(10), val)
}

func Test_read_custom_interface(t *testing.T) {
	should := require.New(t)
	var val MyInterface
	RegisterTypeDecoderFunc("jsoniter.MyInterface", func(ptr unsafe.Pointer, iter *Iterator) {
		*((*MyInterface)(ptr)) = MyString(iter.ReadString())
	})
	err := UnmarshalFromString(`"hello"`, &val)
	should.Nil(err)
	should.Equal("hello", val.Hello())
}

func Test_decode_object_contain_empty_interface(t *testing.T) {
	type TestObject struct {
		Field interface{}
	}
	should := require.New(t)
	obj := TestObject{}
	obj.Field = 1024
	should.Nil(UnmarshalFromString(`{"Field": "hello"}`, &obj))
	should.Equal("hello", obj.Field)
}

func Test_decode_object_contain_non_empty_interface(t *testing.T) {
	type TestObject struct {
		Field MyInterface
	}
	should := require.New(t)
	obj := TestObject{}
	obj.Field = MyString("abc")
	should.Nil(UnmarshalFromString(`{"Field": "hello"}`, &obj))
	should.Equal(MyString("hello"), obj.Field)
}

func Test_encode_object_contain_empty_interface(t *testing.T) {
	type TestObject struct {
		Field interface{}
	}
	should := require.New(t)
	obj := TestObject{}
	obj.Field = 1024
	str, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"Field":1024}`, str)
}

func Test_encode_object_contain_non_empty_interface(t *testing.T) {
	type TestObject struct {
		Field MyInterface
	}
	should := require.New(t)
	obj := TestObject{}
	obj.Field = MyString("hello")
	str, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"Field":"hello"}`, str)
}

func Test_nil_non_empty_interface(t *testing.T) {
	ConfigDefault.(*frozenConfig).cleanEncoders()
	ConfigDefault.(*frozenConfig).cleanDecoders()
	type TestObject struct {
		Field []MyInterface
	}
	should := require.New(t)
	obj := TestObject{}
	b := []byte(`{"Field":["AAA"]}`)
	should.NotNil(json.Unmarshal(b, &obj))
	should.NotNil(Unmarshal(b, &obj))
}

func Test_read_large_number_as_interface(t *testing.T) {
	should := require.New(t)
	var val interface{}
	err := Config{UseNumber: true}.Froze().UnmarshalFromString(`123456789123456789123456789`, &val)
	should.Nil(err)
	output, err := MarshalToString(val)
	should.Nil(err)
	should.Equal(`123456789123456789123456789`, output)
}

func Test_nested_one_field_struct(t *testing.T) {
	should := require.New(t)
	type YetYetAnotherObject struct {
		Field string
	}
	type YetAnotherObject struct {
		Field *YetYetAnotherObject
	}
	type AnotherObject struct {
		Field *YetAnotherObject
	}
	type TestObject struct {
		Me *AnotherObject
	}
	obj := TestObject{&AnotherObject{&YetAnotherObject{&YetYetAnotherObject{"abc"}}}}
	str, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"Me":{"Field":{"Field":{"Field":"abc"}}}}`, str)
	str, err = MarshalToString(&obj)
	should.Nil(err)
	should.Equal(`{"Me":{"Field":{"Field":{"Field":"abc"}}}}`, str)
}

func Test_struct_with_embedded_ptr_with_tag(t *testing.T) {
	type O1 struct {
		O1F string
	}

	type Option struct {
		O1 *O1
	}

	type T struct {
		Option `json:","`
	}
	var obj T
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"O1":null}`, output)
}

func Test_struct_with_one_nil(t *testing.T) {
	type TestObject struct {
		F *float64
	}
	var obj TestObject
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{"F":null}`, output)
}

func Test_struct_with_one_nil_embedded(t *testing.T) {
	type Parent struct {
		Field1 string
		Field2 string
	}
	type TestObject struct {
		*Parent
	}
	obj := TestObject{}
	should := require.New(t)
	bytes, err := json.Marshal(obj)
	should.Nil(err)
	should.Equal("{}", string(bytes))
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`{}`, output)
}

func Test_struct_with_not_nil_embedded(t *testing.T) {
	type Parent struct {
		Field0 string
		Field1 []string
		Field2 map[string]interface{}
	}
	type TestObject struct {
		*Parent
	}
	should := require.New(t)
	var obj TestObject
	err := UnmarshalFromString(`{"Field0":"1","Field1":null,"Field2":{"K":"V"}}`, &obj)
	should.Nil(err)
	should.Nil(obj.Field1)
	should.Equal(map[string]interface{}{"K": "V"}, obj.Field2)
	should.Equal("1", obj.Field0)
}

func Test_array_with_one_nil_ptr(t *testing.T) {
	obj := [1]*float64{nil}
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`[null]`, output)
}

func Test_array_with_one_not_nil_ptr(t *testing.T) {
	two := float64(2)
	obj := [1]*float64{&two}
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`[2]`, output)
}

func Test_embedded_array_with_one_nil(t *testing.T) {
	type TestObject struct {
		Field1 int
		Field2 [1]*float64
	}
	var obj TestObject
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Contains(output, `"Field2":[null]`)
}

func Test_array_with_nothing(t *testing.T) {
	var obj [2]*float64
	should := require.New(t)
	output, err := MarshalToString(obj)
	should.Nil(err)
	should.Equal(`[null,null]`, output)
}

func Test_unmarshal_ptr_to_interface(t *testing.T) {
	type TestData struct {
		Name string `json:"name"`
	}
	should := require.New(t)
	var obj interface{} = &TestData{}
	err := json.Unmarshal([]byte(`{"name":"value"}`), &obj)
	should.Nil(err)
	should.Equal("&{value}", fmt.Sprintf("%v", obj))
	obj = interface{}(&TestData{})
	err = Unmarshal([]byte(`{"name":"value"}`), &obj)
	should.Nil(err)
	should.Equal("&{value}", fmt.Sprintf("%v", obj))
}

func Test_nil_out_null_interface(t *testing.T) {
	type TestData struct {
		Field interface{} `json:"field"`
	}
	should := require.New(t)

	var boolVar bool
	obj := TestData{
		Field: &boolVar,
	}

	data1 := []byte(`{"field": true}`)

	err := Unmarshal(data1, &obj)
	should.Equal(nil, err)
	should.Equal(true, *(obj.Field.(*bool)))

	data2 := []byte(`{"field": null}`)

	err = Unmarshal(data2, &obj)
	should.Equal(nil, err)
	should.Equal(nil, obj.Field)

	// Checking stdlib behavior matches.
	obj2 := TestData{
		Field: &boolVar,
	}

	err = json.Unmarshal(data1, &obj2)
	should.Equal(nil, err)
	should.Equal(true, *(obj2.Field.(*bool)))

	err = json.Unmarshal(data2, &obj2)
	should.Equal(nil, err)
	should.Equal(nil, obj2.Field)
}

func Test_omitempty_nil_interface(t *testing.T) {
	type TestData struct {
		Field interface{} `json:"field,omitempty"`
	}
	should := require.New(t)

	obj := TestData{
		Field: nil,
	}

	js, err := json.Marshal(obj)
	should.Equal(nil, err)
	should.Equal("{}", string(js))

	str, err := MarshalToString(obj)
	should.Equal(nil, err)
	should.Equal(string(js), str)
}

func Test_omitempty_nil_nonempty_interface(t *testing.T) {
	type TestData struct {
		Field MyInterface `json:"field,omitempty"`
	}
	should := require.New(t)

	obj := TestData{
		Field: nil,
	}

	js, err := json.Marshal(obj)
	should.Equal(nil, err)
	should.Equal("{}", string(js))

	str, err := MarshalToString(obj)
	should.Equal(nil, err)
	should.Equal(string(js), str)

	obj.Field = MyString("hello")
	err = UnmarshalFromString(`{"field":null}`, &obj)
	should.Equal(nil, err)
	should.Equal(nil, obj.Field)
}

func Test_marshal_nil_marshaler_interface(t *testing.T) {
	type TestData struct {
		Field json.Marshaler `json:"field"`
	}
	should := require.New(t)

	obj := TestData{
		Field: nil,
	}

	js, err := json.Marshal(obj)
	should.Equal(nil, err)
	should.Equal(`{"field":null}`, string(js))

	str, err := MarshalToString(obj)
	should.Equal(nil, err)
	should.Equal(string(js), str)
}

func Test_marshal_nil_nonempty_interface(t *testing.T) {
	type TestData struct {
		Field MyInterface `json:"field"`
	}
	should := require.New(t)

	obj := TestData{
		Field: nil,
	}

	js, err := json.Marshal(obj)
	should.Equal(nil, err)
	should.Equal(`{"field":null}`, string(js))

	str, err := MarshalToString(obj)
	should.Equal(nil, err)
	should.Equal(string(js), str)

	obj.Field = MyString("hello")
	err = Unmarshal(js, &obj)
	should.Equal(nil, err)
	should.Equal(nil, obj.Field)
}
<|endoftext|>


<commit_before>// Copyright 2017 HyperHQ Inc.
//
// SPDX-License-Identifier: Apache-2.0
//

package main

import (
	"os/signal"
	"sync"
	"testing"
	"os"

	"github.com/stretchr/testify/assert"
)

func TestNewShim(t *testing.T) {
	agent := testSetup(t)
	defer testTearDown(agent)

	id := "foobar"
	pid, err := agent.addContainer(id)
	assert.Nil(t, err, "%s", err)

	shim, err := newShim(mockSockAddr, id, pid)
	assert.Nil(t, err, "%s", err)
	defer shim.agent.Close()

	_, err = newShim(badMockAddr, id, pid)
	assert.NotNil(t, err, "New shim with wrong socket address should fail")
}

func TestShimOps(t *testing.T) {
	agent := testSetup(t)
	defer testTearDown(agent)

	id := "foobar"
	pid, err := agent.addContainer(id)
	assert.Nil(t, err, "%s", err)

	shim, err := newShim(mockSockAddr, id, pid)
	assert.Nil(t, err, "%s", err)
	defer shim.agent.Close()

	shim.resizeTty(os.Stdin)

	wg := &sync.WaitGroup{}
	shim.proxyStdio(wg)

	sigc := shim.forwardAllSignals()
	defer signal.Stop(sigc)

	shim.monitorTtySize(os.Stdin)

	status, err := shim.wait()
	assert.Nil(t, err, "%s", err)
	assert.Equal(t, status, int32(0), "process fail status %d", status)
}
<commit_msg>TestShimOps: Change os.Stdin to pty.Open<commit_after>// Copyright 2017 HyperHQ Inc.
//
// SPDX-License-Identifier: Apache-2.0
//

package main

import (
	"os/signal"
	"sync"
	"testing"

	"github.com/kr/pty"
	"github.com/stretchr/testify/assert"
)

func TestNewShim(t *testing.T) {
	agent := testSetup(t)
	defer testTearDown(agent)

	id := "foobar"
	pid, err := agent.addContainer(id)
	assert.Nil(t, err, "%s", err)

	shim, err := newShim(mockSockAddr, id, pid)
	assert.Nil(t, err, "%s", err)
	defer shim.agent.Close()

	_, err = newShim(badMockAddr, id, pid)
	assert.NotNil(t, err, "New shim with wrong socket address should fail")
}

func TestShimOps(t *testing.T) {
	agent := testSetup(t)
	defer testTearDown(agent)

	id := "foobar"
	pid, err := agent.addContainer(id)
	assert.Nil(t, err, "%s", err)

	shim, err := newShim(mockSockAddr, id, pid)
	assert.Nil(t, err, "%s", err)
	defer shim.agent.Close()

	_, tty, err := pty.Open()
	assert.Nil(t, err, "%s", err)
	defer tty.Close()

	shim.resizeTty(tty)

	wg := &sync.WaitGroup{}
	shim.proxyStdio(wg)

	sigc := shim.forwardAllSignals()
	defer signal.Stop(sigc)

	shim.monitorTtySize(tty)

	status, err := shim.wait()
	assert.Nil(t, err, "%s", err)
	assert.Equal(t, status, int32(0), "process fail status %d", status)
}
<|endoftext|>


<commit_before>/*
 * MinIO Cloud Storage, (C) 2016, 2017, 2018 MinIO, Inc.
 *
 * 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 cmd

import (
	"encoding/json"
	"encoding/xml"
	"errors"
	"io"
	"net/http"
	"net/url"
	"path"
	"reflect"
	"time"

	"github.com/gorilla/mux"
	xhttp "github.com/minio/minio/cmd/http"
	"github.com/minio/minio/cmd/logger"
	"github.com/minio/minio/pkg/event"
	xnet "github.com/minio/minio/pkg/net"
	"github.com/minio/minio/pkg/policy"
)

const (
	bucketConfigPrefix       = "buckets"
	bucketNotificationConfig = "notification.xml"
)

var errNoSuchNotifications = errors.New("The specified bucket does not have bucket notifications")

// GetBucketNotificationHandler - This HTTP handler returns event notification configuration
// as per http://docs.aws.amazon.com/AmazonS3/latest/dev/NotificationHowTo.html.
// It returns empty configuration if its not set.
func (api objectAPIHandlers) GetBucketNotificationHandler(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "GetBucketNotification")

	defer logger.AuditLog(w, r, "GetBucketNotification", mustGetClaimsFromToken(r))

	vars := mux.Vars(r)
	bucketName := vars["bucket"]

	objAPI := api.ObjectAPI()
	if objAPI == nil {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrServerNotInitialized), r.URL, guessIsBrowserReq(r))
		return
	}

	if !objAPI.IsNotificationSupported() {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrNotImplemented), r.URL, guessIsBrowserReq(r))
		return
	}

	if s3Error := checkRequestAuthType(ctx, r, policy.GetBucketNotificationAction, bucketName, ""); s3Error != ErrNone {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(s3Error), r.URL, guessIsBrowserReq(r))
		return
	}

	_, err := objAPI.GetBucketInfo(ctx, bucketName)
	if err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	// Construct path to notification.xml for the given bucket.
	configFile := path.Join(bucketConfigPrefix, bucketName, bucketNotificationConfig)

	var config = event.Config{}
	configData, err := readConfig(ctx, objAPI, configFile)
	if err != nil {
		if err != errConfigNotFound {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}
		config.XMLNS = "http://s3.amazonaws.com/doc/2006-03-01/"
		config.SetRegion(globalServerRegion)
		notificationBytes, err := xml.Marshal(config)
		if err != nil {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}

		writeSuccessResponseXML(w, notificationBytes)
		return
	}

	config.SetRegion(globalServerRegion)

	if err = xml.Unmarshal(configData, &config); err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	if err = config.Validate(globalServerRegion, globalNotificationSys.targetList); err != nil {
		arnErr, ok := err.(*event.ErrARNNotFound)
		if !ok {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}
		for i, queue := range config.QueueList {
			// Remove ARN not found queues, because we previously allowed
			// adding unexpected entries into the config.
			//
			// With newer config disallowing changing / turning off
			// notification targets without removing ARN in notification
			// configuration we won't see this problem anymore.
			if reflect.DeepEqual(queue.ARN, arnErr.ARN) {
				config.QueueList = append(config.QueueList[:i],
					config.QueueList[i+1:]...)
			}
			// This is a one time activity we shall do this
			// here and allow stale ARN to be removed. We shall
			// never reach a stage where we will have stale
			// notification configs.
		}
	}

	// If xml namespace is empty, set a default value before returning.
	if config.XMLNS == "" {
		config.XMLNS = "http://s3.amazonaws.com/doc/2006-03-01/"
	}

	notificationBytes, err := xml.Marshal(config)
	if err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	writeSuccessResponseXML(w, notificationBytes)
}

// PutBucketNotificationHandler - This HTTP handler stores given notification configuration as per
// http://docs.aws.amazon.com/AmazonS3/latest/dev/NotificationHowTo.html.
func (api objectAPIHandlers) PutBucketNotificationHandler(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "PutBucketNotification")

	defer logger.AuditLog(w, r, "PutBucketNotification", mustGetClaimsFromToken(r))

	objectAPI := api.ObjectAPI()
	if objectAPI == nil {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrServerNotInitialized), r.URL, guessIsBrowserReq(r))
		return
	}

	if !objectAPI.IsNotificationSupported() {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrNotImplemented), r.URL, guessIsBrowserReq(r))
		return
	}

	vars := mux.Vars(r)
	bucketName := vars["bucket"]

	if s3Error := checkRequestAuthType(ctx, r, policy.PutBucketNotificationAction, bucketName, ""); s3Error != ErrNone {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(s3Error), r.URL, guessIsBrowserReq(r))
		return
	}

	_, err := objectAPI.GetBucketInfo(ctx, bucketName)
	if err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	// PutBucketNotification always needs a Content-Length.
	if r.ContentLength <= 0 {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrMissingContentLength), r.URL, guessIsBrowserReq(r))
		return
	}

	lreader := io.LimitReader(r.Body, r.ContentLength)
	config, err := event.ParseConfig(lreader, globalServerRegion, globalNotificationSys.targetList)
	if err != nil {
		apiErr := errorCodes.ToAPIErr(ErrMalformedXML)
		if event.IsEventError(err) {
			apiErr = toAPIError(ctx, err)
		}
		writeErrorResponse(ctx, w, apiErr, r.URL, guessIsBrowserReq(r))
		return
	}

	if err = saveNotificationConfig(ctx, objectAPI, bucketName, config); err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	rulesMap := config.ToRulesMap()
	globalNotificationSys.AddRulesMap(bucketName, rulesMap)
	globalNotificationSys.PutBucketNotification(ctx, bucketName, rulesMap)

	writeSuccessResponseHeadersOnly(w)
}

func (api objectAPIHandlers) ListenBucketNotificationHandler(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "ListenBucketNotification")

	defer logger.AuditLog(w, r, "ListenBucketNotification", mustGetClaimsFromToken(r))

	// Validate if bucket exists.
	objAPI := api.ObjectAPI()
	if objAPI == nil {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrServerNotInitialized), r.URL, guessIsBrowserReq(r))
		return
	}

	if !objAPI.IsNotificationSupported() {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrNotImplemented), r.URL, guessIsBrowserReq(r))
		return
	}

	if !objAPI.IsListenBucketSupported() {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrNotImplemented), r.URL, guessIsBrowserReq(r))
		return
	}

	vars := mux.Vars(r)
	bucketName := vars["bucket"]

	values := r.URL.Query()

	var prefix string
	if len(values["prefix"]) > 1 {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrFilterNamePrefix), r.URL, guessIsBrowserReq(r))
		return
	}

	if len(values["prefix"]) == 1 {
		if err := event.ValidateFilterRuleValue(values["prefix"][0]); err != nil {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}

		prefix = values["prefix"][0]
	}

	var suffix string
	if len(values["suffix"]) > 1 {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrFilterNameSuffix), r.URL, guessIsBrowserReq(r))
		return
	}

	if len(values["suffix"]) == 1 {
		if err := event.ValidateFilterRuleValue(values["suffix"][0]); err != nil {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}

		suffix = values["suffix"][0]
	}

	pattern := event.NewPattern(prefix, suffix)

	eventNames := []event.Name{}
	for _, s := range values["events"] {
		eventName, err := event.ParseName(s)
		if err != nil {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}

		eventNames = append(eventNames, eventName)
	}

	if _, err := objAPI.GetBucketInfo(ctx, bucketName); err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	host, err := xnet.ParseHost(r.RemoteAddr)
	if err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	rulesMap := event.NewRulesMap(eventNames, pattern,
		event.TargetID{ID: "listen" + "+" + mustGetUUID() + "+" + host.Name, Name: host.Port.String()})

	w.Header().Set(xhttp.ContentType, "text/event-stream")

	doneCh := make(chan struct{})
	defer close(doneCh)

	// Listen Publisher and peer-listen-client uses nonblocking send and hence does not wait for slow receivers.
	// Use buffered channel to take care of burst sends or slow w.Write()
	listenCh := make(chan interface{}, 4000)

	peers := getRestClients(globalEndpoints)

	globalHTTPListen.Subscribe(listenCh, doneCh, func(evI interface{}) bool {
		ev, ok := evI.(event.Event)
		if !ok {
			return false
		}
		objectName, uerr := url.QueryUnescape(ev.S3.Object.Key)
		if uerr != nil {
			return false
		}
		return len(rulesMap.Match(ev.EventName, objectName).ToSlice()) != 0
	})

	for _, peer := range peers {
		if peer == nil {
			continue
		}
		peer.Listen(listenCh, doneCh)
	}

	keepAliveTicker := time.NewTicker(500 * time.Millisecond)
	defer keepAliveTicker.Stop()

	enc := json.NewEncoder(w)
	for {
		select {
		case evI := <-listenCh:
			ev := evI.(event.Event)
			if err := enc.Encode(struct{ Records []event.Event }{[]event.Event{ev}}); err != nil {
				return
			}
			w.(http.Flusher).Flush()
		case <-keepAliveTicker.C:
			if _, err := w.Write([]byte(" ")); err != nil {
				return
			}
			w.(http.Flusher).Flush()
		case <-GlobalServiceDoneCh:
			return
		}
	}

}
<commit_msg>Always use SourceIP for host target filtering (#8649)<commit_after>/*
 * MinIO Cloud Storage, (C) 2016, 2017, 2018 MinIO, Inc.
 *
 * 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 cmd

import (
	"encoding/json"
	"encoding/xml"
	"errors"
	"io"
	"net/http"
	"net/url"
	"path"
	"reflect"
	"time"

	"github.com/gorilla/mux"
	xhttp "github.com/minio/minio/cmd/http"
	"github.com/minio/minio/cmd/logger"
	"github.com/minio/minio/pkg/event"
	"github.com/minio/minio/pkg/handlers"
	xnet "github.com/minio/minio/pkg/net"
	"github.com/minio/minio/pkg/policy"
)

const (
	bucketConfigPrefix       = "buckets"
	bucketNotificationConfig = "notification.xml"
)

var errNoSuchNotifications = errors.New("The specified bucket does not have bucket notifications")

// GetBucketNotificationHandler - This HTTP handler returns event notification configuration
// as per http://docs.aws.amazon.com/AmazonS3/latest/dev/NotificationHowTo.html.
// It returns empty configuration if its not set.
func (api objectAPIHandlers) GetBucketNotificationHandler(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "GetBucketNotification")

	defer logger.AuditLog(w, r, "GetBucketNotification", mustGetClaimsFromToken(r))

	vars := mux.Vars(r)
	bucketName := vars["bucket"]

	objAPI := api.ObjectAPI()
	if objAPI == nil {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrServerNotInitialized), r.URL, guessIsBrowserReq(r))
		return
	}

	if !objAPI.IsNotificationSupported() {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrNotImplemented), r.URL, guessIsBrowserReq(r))
		return
	}

	if s3Error := checkRequestAuthType(ctx, r, policy.GetBucketNotificationAction, bucketName, ""); s3Error != ErrNone {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(s3Error), r.URL, guessIsBrowserReq(r))
		return
	}

	_, err := objAPI.GetBucketInfo(ctx, bucketName)
	if err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	// Construct path to notification.xml for the given bucket.
	configFile := path.Join(bucketConfigPrefix, bucketName, bucketNotificationConfig)

	var config = event.Config{}
	configData, err := readConfig(ctx, objAPI, configFile)
	if err != nil {
		if err != errConfigNotFound {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}
		config.XMLNS = "http://s3.amazonaws.com/doc/2006-03-01/"
		config.SetRegion(globalServerRegion)
		notificationBytes, err := xml.Marshal(config)
		if err != nil {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}

		writeSuccessResponseXML(w, notificationBytes)
		return
	}

	config.SetRegion(globalServerRegion)

	if err = xml.Unmarshal(configData, &config); err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	if err = config.Validate(globalServerRegion, globalNotificationSys.targetList); err != nil {
		arnErr, ok := err.(*event.ErrARNNotFound)
		if !ok {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}
		for i, queue := range config.QueueList {
			// Remove ARN not found queues, because we previously allowed
			// adding unexpected entries into the config.
			//
			// With newer config disallowing changing / turning off
			// notification targets without removing ARN in notification
			// configuration we won't see this problem anymore.
			if reflect.DeepEqual(queue.ARN, arnErr.ARN) {
				config.QueueList = append(config.QueueList[:i],
					config.QueueList[i+1:]...)
			}
			// This is a one time activity we shall do this
			// here and allow stale ARN to be removed. We shall
			// never reach a stage where we will have stale
			// notification configs.
		}
	}

	// If xml namespace is empty, set a default value before returning.
	if config.XMLNS == "" {
		config.XMLNS = "http://s3.amazonaws.com/doc/2006-03-01/"
	}

	notificationBytes, err := xml.Marshal(config)
	if err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	writeSuccessResponseXML(w, notificationBytes)
}

// PutBucketNotificationHandler - This HTTP handler stores given notification configuration as per
// http://docs.aws.amazon.com/AmazonS3/latest/dev/NotificationHowTo.html.
func (api objectAPIHandlers) PutBucketNotificationHandler(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "PutBucketNotification")

	defer logger.AuditLog(w, r, "PutBucketNotification", mustGetClaimsFromToken(r))

	objectAPI := api.ObjectAPI()
	if objectAPI == nil {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrServerNotInitialized), r.URL, guessIsBrowserReq(r))
		return
	}

	if !objectAPI.IsNotificationSupported() {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrNotImplemented), r.URL, guessIsBrowserReq(r))
		return
	}

	vars := mux.Vars(r)
	bucketName := vars["bucket"]

	if s3Error := checkRequestAuthType(ctx, r, policy.PutBucketNotificationAction, bucketName, ""); s3Error != ErrNone {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(s3Error), r.URL, guessIsBrowserReq(r))
		return
	}

	_, err := objectAPI.GetBucketInfo(ctx, bucketName)
	if err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	// PutBucketNotification always needs a Content-Length.
	if r.ContentLength <= 0 {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrMissingContentLength), r.URL, guessIsBrowserReq(r))
		return
	}

	lreader := io.LimitReader(r.Body, r.ContentLength)
	config, err := event.ParseConfig(lreader, globalServerRegion, globalNotificationSys.targetList)
	if err != nil {
		apiErr := errorCodes.ToAPIErr(ErrMalformedXML)
		if event.IsEventError(err) {
			apiErr = toAPIError(ctx, err)
		}
		writeErrorResponse(ctx, w, apiErr, r.URL, guessIsBrowserReq(r))
		return
	}

	if err = saveNotificationConfig(ctx, objectAPI, bucketName, config); err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	rulesMap := config.ToRulesMap()
	globalNotificationSys.AddRulesMap(bucketName, rulesMap)
	globalNotificationSys.PutBucketNotification(ctx, bucketName, rulesMap)

	writeSuccessResponseHeadersOnly(w)
}

func (api objectAPIHandlers) ListenBucketNotificationHandler(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "ListenBucketNotification")

	defer logger.AuditLog(w, r, "ListenBucketNotification", mustGetClaimsFromToken(r))

	// Validate if bucket exists.
	objAPI := api.ObjectAPI()
	if objAPI == nil {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrServerNotInitialized), r.URL, guessIsBrowserReq(r))
		return
	}

	if !objAPI.IsNotificationSupported() {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrNotImplemented), r.URL, guessIsBrowserReq(r))
		return
	}

	if !objAPI.IsListenBucketSupported() {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrNotImplemented), r.URL, guessIsBrowserReq(r))
		return
	}

	vars := mux.Vars(r)
	bucketName := vars["bucket"]

	values := r.URL.Query()

	var prefix string
	if len(values["prefix"]) > 1 {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrFilterNamePrefix), r.URL, guessIsBrowserReq(r))
		return
	}

	if len(values["prefix"]) == 1 {
		if err := event.ValidateFilterRuleValue(values["prefix"][0]); err != nil {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}

		prefix = values["prefix"][0]
	}

	var suffix string
	if len(values["suffix"]) > 1 {
		writeErrorResponse(ctx, w, errorCodes.ToAPIErr(ErrFilterNameSuffix), r.URL, guessIsBrowserReq(r))
		return
	}

	if len(values["suffix"]) == 1 {
		if err := event.ValidateFilterRuleValue(values["suffix"][0]); err != nil {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}

		suffix = values["suffix"][0]
	}

	pattern := event.NewPattern(prefix, suffix)

	eventNames := []event.Name{}
	for _, s := range values["events"] {
		eventName, err := event.ParseName(s)
		if err != nil {
			writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
			return
		}

		eventNames = append(eventNames, eventName)
	}

	if _, err := objAPI.GetBucketInfo(ctx, bucketName); err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	host, err := xnet.ParseHost(handlers.GetSourceIP(r))
	if err != nil {
		writeErrorResponse(ctx, w, toAPIError(ctx, err), r.URL, guessIsBrowserReq(r))
		return
	}

	rulesMap := event.NewRulesMap(eventNames, pattern,
		event.TargetID{ID: "listen" + "+" + mustGetUUID() + "+" + host.Name, Name: host.Port.String()})

	w.Header().Set(xhttp.ContentType, "text/event-stream")

	doneCh := make(chan struct{})
	defer close(doneCh)

	// Listen Publisher and peer-listen-client uses nonblocking send and hence does not wait for slow receivers.
	// Use buffered channel to take care of burst sends or slow w.Write()
	listenCh := make(chan interface{}, 4000)

	peers := getRestClients(globalEndpoints)

	globalHTTPListen.Subscribe(listenCh, doneCh, func(evI interface{}) bool {
		ev, ok := evI.(event.Event)
		if !ok {
			return false
		}
		objectName, uerr := url.QueryUnescape(ev.S3.Object.Key)
		if uerr != nil {
			return false
		}
		return len(rulesMap.Match(ev.EventName, objectName).ToSlice()) != 0
	})

	for _, peer := range peers {
		if peer == nil {
			continue
		}
		peer.Listen(listenCh, doneCh)
	}

	keepAliveTicker := time.NewTicker(500 * time.Millisecond)
	defer keepAliveTicker.Stop()

	enc := json.NewEncoder(w)
	for {
		select {
		case evI := <-listenCh:
			ev := evI.(event.Event)
			if err := enc.Encode(struct{ Records []event.Event }{[]event.Event{ev}}); err != nil {
				return
			}
			w.(http.Flusher).Flush()
		case <-keepAliveTicker.C:
			if _, err := w.Write([]byte(" ")); err != nil {
				return
			}
			w.(http.Flusher).Flush()
		case <-GlobalServiceDoneCh:
			return
		}
	}

}
<|endoftext|>