<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strings"
)

type proxyHandler struct{}

func (h proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Host == "203.0.113.1" {
		http.DefaultServeMux.ServeHTTP(w, r)
		return
	}

	r.Header.Add("Via", "Redwood")
	client := r.RemoteAddr
	host, _, err := net.SplitHostPort(client)
	if err == nil {
		client = host
	}
	r.Header.Add("X-Forwarded-For", client)
	r.Header.Del("Accept-Encoding")

	// Reconstruct the URL if this is a transparent proxy.
	if r.URL.Host == "" {
		r.URL.Host = r.Host
		if r.URL.Scheme == "" {
			r.URL.Scheme = "http"
			// TODO: handle HTTPS
		}
	}

	sc := scorecard{
		tally: URLRules.MatchingRules(r.URL),
	}
	sc.calculate(client)
	if sc.action == BLOCK {
		showBlockPage(w, r, &sc)
		return
	}

	changeQuery(r.URL)

	resp, err := http.DefaultTransport.RoundTrip(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}

	contentType, action := checkContentType(resp)

	switch action {
	case BLOCK:
		sc.action = BLOCK
		sc.blocked = []string{"blocked-mime"}
		showBlockPage(w, r, &sc)
		return

	case ALLOW:
		copyResponseHeader(w, resp)
		io.Copy(w, resp.Body)
		return
	}

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(fmt.Errorf("error while reading response body: %s", err))
	}

	modified := false
	charset := findCharset(resp.Header.Get("Content-Type"), content)
	if strings.Contains(contentType, "html") {
		modified = pruneContent(r.URL, &content, charset)
		if modified {
			resp.Header.Set("Content-Type", "text/html; charset=utf-8")
			charset = "utf-8"
		}
	}

	scanContent(content, contentType, charset, sc.tally)
	sc.calculate(client)

	if sc.action == BLOCK {
		showBlockPage(w, r, &sc)
		return
	}

	copyResponseHeader(w, resp)
	w.Write(content)
}

// copyResponseHeader writes resp's header and status code to w.
func copyResponseHeader(w http.ResponseWriter, resp *http.Response) {
	newHeader := w.Header()
	for key, values := range resp.Header {
		for _, v := range values {
			newHeader.Add(key, v)
		}
	}

	w.WriteHeader(resp.StatusCode)
}
<commit_msg>Make sure response body gets closed.<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strings"
)

type proxyHandler struct{}

func (h proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Host == "203.0.113.1" {
		http.DefaultServeMux.ServeHTTP(w, r)
		return
	}

	r.Header.Add("Via", "Redwood")
	client := r.RemoteAddr
	host, _, err := net.SplitHostPort(client)
	if err == nil {
		client = host
	}
	r.Header.Add("X-Forwarded-For", client)
	r.Header.Del("Accept-Encoding")

	// Reconstruct the URL if this is a transparent proxy.
	if r.URL.Host == "" {
		r.URL.Host = r.Host
		if r.URL.Scheme == "" {
			r.URL.Scheme = "http"
			// TODO: handle HTTPS
		}
	}

	sc := scorecard{
		tally: URLRules.MatchingRules(r.URL),
	}
	sc.calculate(client)
	if sc.action == BLOCK {
		showBlockPage(w, r, &sc)
		return
	}

	changeQuery(r.URL)

	resp, err := http.DefaultTransport.RoundTrip(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer resp.Body.Close()

	contentType, action := checkContentType(resp)

	switch action {
	case BLOCK:
		sc.action = BLOCK
		sc.blocked = []string{"blocked-mime"}
		showBlockPage(w, r, &sc)
		return

	case ALLOW:
		copyResponseHeader(w, resp)
		io.Copy(w, resp.Body)
		return
	}

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(fmt.Errorf("error while reading response body: %s", err))
	}

	modified := false
	charset := findCharset(resp.Header.Get("Content-Type"), content)
	if strings.Contains(contentType, "html") {
		modified = pruneContent(r.URL, &content, charset)
		if modified {
			resp.Header.Set("Content-Type", "text/html; charset=utf-8")
			charset = "utf-8"
		}
	}

	scanContent(content, contentType, charset, sc.tally)
	sc.calculate(client)

	if sc.action == BLOCK {
		showBlockPage(w, r, &sc)
		return
	}

	copyResponseHeader(w, resp)
	w.Write(content)
}

// copyResponseHeader writes resp's header and status code to w.
func copyResponseHeader(w http.ResponseWriter, resp *http.Response) {
	newHeader := w.Header()
	for key, values := range resp.Header {
		for _, v := range values {
			newHeader.Add(key, v)
		}
	}

	w.WriteHeader(resp.StatusCode)
}
<|endoftext|>


<commit_before>package pgmock

import (
	"encoding/json"
	"fmt"
	"net"

	"github.com/jackc/pgx/pgproto3"
)

type Proxy struct {
	frontend *pgproto3.Frontend
	backend  *pgproto3.Backend

	frontendConn net.Conn
	backendConn  net.Conn
}

func NewProxy(frontendConn, backendConn net.Conn) (*Proxy, error) {
	backend, err := pgproto3.NewBackend(frontendConn, frontendConn)
	if err != nil {
		return nil, err
	}
	frontend, err := pgproto3.NewFrontend(backendConn, backendConn)
	if err != nil {
		return nil, err
	}

	proxy := &Proxy{
		frontend: frontend,
		backend:  backend,

		frontendConn: frontendConn,
		backendConn:  backendConn,
	}

	return proxy, nil
}

func (p *Proxy) Run() error {
	defer p.Close()

	frontendErrChan := make(chan error, 1)
	frontendMsgChan := make(chan pgproto3.FrontendMessage)
	go p.readClientConn(frontendMsgChan, frontendErrChan)

	backendErrChan := make(chan error, 1)
	backendMsgChan := make(chan pgproto3.BackendMessage)
	go p.readServerConn(backendMsgChan, backendErrChan)

	for {
		select {
		case msg := <-frontendMsgChan:
			fmt.Print("frontend: ")
			buf, err := json.Marshal(msg)
			if err != nil {
				return err
			}
			fmt.Println(string(buf))

			err = p.frontend.Send(msg)
			if err != nil {
				return err
			}
		case msg := <-backendMsgChan:
			fmt.Print("backend: ")
			buf, err := json.Marshal(msg)
			if err != nil {
				return err
			}
			fmt.Println(string(buf))

			err = p.backend.Send(msg)
			if err != nil {
				return err
			}
		case err := <-frontendErrChan:
			return err
		case err := <-backendErrChan:
			return err
		}
	}
}

func (p *Proxy) Close() error {
	frontendCloseErr := p.frontendConn.Close()
	backendCloseErr := p.backendConn.Close()

	if frontendCloseErr != nil {
		return frontendCloseErr
	}
	return backendCloseErr
}

func (p *Proxy) readClientConn(msgChan chan pgproto3.FrontendMessage, errChan chan error) {
	startupMessage, err := p.backend.ReceiveStartupMessage()
	if err != nil {
		errChan <- err
		return
	}

	msgChan <- startupMessage

	for {
		msg, err := p.backend.Receive()
		if err != nil {
			errChan <- err
			return
		}

		msgChan <- msg
	}
}

func (p *Proxy) readServerConn(msgChan chan pgproto3.BackendMessage, errChan chan error) {
	for {
		msg, err := p.frontend.Receive()
		if err != nil {
			errChan <- err
			return
		}

		msgChan <- msg
	}
}
<commit_msg>Fix race condition with flyweight messages<commit_after>package pgmock

import (
	"encoding/json"
	"fmt"
	"net"

	"github.com/jackc/pgx/pgproto3"
)

type Proxy struct {
	frontend *pgproto3.Frontend
	backend  *pgproto3.Backend

	frontendConn net.Conn
	backendConn  net.Conn
}

func NewProxy(frontendConn, backendConn net.Conn) (*Proxy, error) {
	backend, err := pgproto3.NewBackend(frontendConn, frontendConn)
	if err != nil {
		return nil, err
	}
	frontend, err := pgproto3.NewFrontend(backendConn, backendConn)
	if err != nil {
		return nil, err
	}

	proxy := &Proxy{
		frontend: frontend,
		backend:  backend,

		frontendConn: frontendConn,
		backendConn:  backendConn,
	}

	return proxy, nil
}

func (p *Proxy) Run() error {
	defer p.Close()

	frontendErrChan := make(chan error, 1)
	frontendMsgChan := make(chan pgproto3.FrontendMessage)
	frontendNextChan := make(chan struct{})
	go p.readClientConn(frontendMsgChan, frontendNextChan, frontendErrChan)

	backendErrChan := make(chan error, 1)
	backendMsgChan := make(chan pgproto3.BackendMessage)
	backendNextChan := make(chan struct{})
	go p.readServerConn(backendMsgChan, backendNextChan, backendErrChan)

	for {
		select {
		case msg := <-frontendMsgChan:
			fmt.Print("frontend: ")
			buf, err := json.Marshal(msg)
			if err != nil {
				return err
			}
			fmt.Println(string(buf))

			err = p.frontend.Send(msg)
			if err != nil {
				return err
			}
			frontendNextChan <- struct{}{}
		case msg := <-backendMsgChan:
			fmt.Print("backend: ")
			buf, err := json.Marshal(msg)
			if err != nil {
				return err
			}
			fmt.Println(string(buf))

			err = p.backend.Send(msg)
			if err != nil {
				return err
			}
			backendNextChan <- struct{}{}
		case err := <-frontendErrChan:
			return err
		case err := <-backendErrChan:
			return err
		}
	}
}

func (p *Proxy) Close() error {
	frontendCloseErr := p.frontendConn.Close()
	backendCloseErr := p.backendConn.Close()

	if frontendCloseErr != nil {
		return frontendCloseErr
	}
	return backendCloseErr
}

func (p *Proxy) readClientConn(msgChan chan pgproto3.FrontendMessage, nextChan chan struct{}, errChan chan error) {
	startupMessage, err := p.backend.ReceiveStartupMessage()
	if err != nil {
		errChan <- err
		return
	}

	msgChan <- startupMessage
	<-nextChan

	for {
		msg, err := p.backend.Receive()
		if err != nil {
			errChan <- err
			return
		}

		msgChan <- msg
		<-nextChan
	}
}

func (p *Proxy) readServerConn(msgChan chan pgproto3.BackendMessage, nextChan chan struct{}, errChan chan error) {
	for {
		msg, err := p.frontend.Receive()
		if err != nil {
			errChan <- err
			return
		}

		msgChan <- msg

		<-nextChan
	}
}
<|endoftext|>


<commit_before>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cipher_test

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"testing"
)

// cfbTests contains the test vectors from
// https://csrc.nist.gov/publications/nistpubs/800-38a/sp800-38a.pdf, section
// F.3.13.
var cfbTests = []struct {
	key, iv, plaintext, ciphertext string
}{
	{
		"2b7e151628aed2a6abf7158809cf4f3c",
		"000102030405060708090a0b0c0d0e0f",
		"6bc1bee22e409f96e93d7e117393172a",
		"3b3fd92eb72dad20333449f8e83cfb4a",
	},
	{
		"2b7e151628aed2a6abf7158809cf4f3c",
		"3B3FD92EB72DAD20333449F8E83CFB4A",
		"ae2d8a571e03ac9c9eb76fac45af8e51",
		"c8a64537a0b3a93fcde3cdad9f1ce58b",
	},
	{
		"2b7e151628aed2a6abf7158809cf4f3c",
		"C8A64537A0B3A93FCDE3CDAD9F1CE58B",
		"30c81c46a35ce411e5fbc1191a0a52ef",
		"26751f67a3cbb140b1808cf187a4f4df",
	},
	{
		"2b7e151628aed2a6abf7158809cf4f3c",
		"26751F67A3CBB140B1808CF187A4F4DF",
		"f69f2445df4f9b17ad2b417be66c3710",
		"c04b05357c5d1c0eeac4c66f9ff7f2e6",
	},
}

func TestCFBVectors(t *testing.T) {
	for i, test := range cfbTests {
		key, err := hex.DecodeString(test.key)
		if err != nil {
			t.Fatal(err)
		}
		iv, err := hex.DecodeString(test.iv)
		if err != nil {
			t.Fatal(err)
		}
		plaintext, err := hex.DecodeString(test.plaintext)
		if err != nil {
			t.Fatal(err)
		}
		expected, err := hex.DecodeString(test.ciphertext)
		if err != nil {
			t.Fatal(err)
		}

		block, err := aes.NewCipher(key)
		if err != nil {
			t.Fatal(err)
		}

		ciphertext := make([]byte, len(plaintext))
		cfb := cipher.NewCFBEncrypter(block, iv)
		cfb.XORKeyStream(ciphertext, plaintext)

		if !bytes.Equal(ciphertext, expected) {
			t.Errorf("#%d: wrong output: got %x, expected %x", i, ciphertext, expected)
		}

		cfbdec := cipher.NewCFBDecrypter(block, iv)
		plaintextCopy := make([]byte, len(ciphertext))
		cfbdec.XORKeyStream(plaintextCopy, ciphertext)

		if !bytes.Equal(plaintextCopy, plaintextCopy) {
			t.Errorf("#%d: wrong plaintext: got %x, expected %x", i, plaintextCopy, plaintext)
		}
	}
}

func TestCFBInverse(t *testing.T) {
	block, err := aes.NewCipher(commonKey128)
	if err != nil {
		t.Error(err)
		return
	}

	plaintext := []byte("this is the plaintext. this is the plaintext.")
	iv := make([]byte, block.BlockSize())
	rand.Reader.Read(iv)
	cfb := cipher.NewCFBEncrypter(block, iv)
	ciphertext := make([]byte, len(plaintext))
	copy(ciphertext, plaintext)
	cfb.XORKeyStream(ciphertext, ciphertext)

	cfbdec := cipher.NewCFBDecrypter(block, iv)
	plaintextCopy := make([]byte, len(plaintext))
	copy(plaintextCopy, ciphertext)
	cfbdec.XORKeyStream(plaintextCopy, plaintextCopy)

	if !bytes.Equal(plaintextCopy, plaintext) {
		t.Errorf("got: %x, want: %x", plaintextCopy, plaintext)
	}
}
<commit_msg>crypto/cipher: fix duplicated arguments to bytes.Equal in test<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cipher_test

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"testing"
)

// cfbTests contains the test vectors from
// https://csrc.nist.gov/publications/nistpubs/800-38a/sp800-38a.pdf, section
// F.3.13.
var cfbTests = []struct {
	key, iv, plaintext, ciphertext string
}{
	{
		"2b7e151628aed2a6abf7158809cf4f3c",
		"000102030405060708090a0b0c0d0e0f",
		"6bc1bee22e409f96e93d7e117393172a",
		"3b3fd92eb72dad20333449f8e83cfb4a",
	},
	{
		"2b7e151628aed2a6abf7158809cf4f3c",
		"3B3FD92EB72DAD20333449F8E83CFB4A",
		"ae2d8a571e03ac9c9eb76fac45af8e51",
		"c8a64537a0b3a93fcde3cdad9f1ce58b",
	},
	{
		"2b7e151628aed2a6abf7158809cf4f3c",
		"C8A64537A0B3A93FCDE3CDAD9F1CE58B",
		"30c81c46a35ce411e5fbc1191a0a52ef",
		"26751f67a3cbb140b1808cf187a4f4df",
	},
	{
		"2b7e151628aed2a6abf7158809cf4f3c",
		"26751F67A3CBB140B1808CF187A4F4DF",
		"f69f2445df4f9b17ad2b417be66c3710",
		"c04b05357c5d1c0eeac4c66f9ff7f2e6",
	},
}

func TestCFBVectors(t *testing.T) {
	for i, test := range cfbTests {
		key, err := hex.DecodeString(test.key)
		if err != nil {
			t.Fatal(err)
		}
		iv, err := hex.DecodeString(test.iv)
		if err != nil {
			t.Fatal(err)
		}
		plaintext, err := hex.DecodeString(test.plaintext)
		if err != nil {
			t.Fatal(err)
		}
		expected, err := hex.DecodeString(test.ciphertext)
		if err != nil {
			t.Fatal(err)
		}

		block, err := aes.NewCipher(key)
		if err != nil {
			t.Fatal(err)
		}

		ciphertext := make([]byte, len(plaintext))
		cfb := cipher.NewCFBEncrypter(block, iv)
		cfb.XORKeyStream(ciphertext, plaintext)

		if !bytes.Equal(ciphertext, expected) {
			t.Errorf("#%d: wrong output: got %x, expected %x", i, ciphertext, expected)
		}

		cfbdec := cipher.NewCFBDecrypter(block, iv)
		plaintextCopy := make([]byte, len(ciphertext))
		cfbdec.XORKeyStream(plaintextCopy, ciphertext)

		if !bytes.Equal(plaintextCopy, plaintext) {
			t.Errorf("#%d: wrong plaintext: got %x, expected %x", i, plaintextCopy, plaintext)
		}
	}
}

func TestCFBInverse(t *testing.T) {
	block, err := aes.NewCipher(commonKey128)
	if err != nil {
		t.Error(err)
		return
	}

	plaintext := []byte("this is the plaintext. this is the plaintext.")
	iv := make([]byte, block.BlockSize())
	rand.Reader.Read(iv)
	cfb := cipher.NewCFBEncrypter(block, iv)
	ciphertext := make([]byte, len(plaintext))
	copy(ciphertext, plaintext)
	cfb.XORKeyStream(ciphertext, ciphertext)

	cfbdec := cipher.NewCFBDecrypter(block, iv)
	plaintextCopy := make([]byte, len(plaintext))
	copy(plaintextCopy, ciphertext)
	cfbdec.XORKeyStream(plaintextCopy, plaintextCopy)

	if !bytes.Equal(plaintextCopy, plaintext) {
		t.Errorf("got: %x, want: %x", plaintextCopy, plaintext)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
)

var port = flag.String("p", "8888", "Port to listen on")

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage:	%v -p <port> <storage-path>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func errorResponse(w http.ResponseWriter, e error) {
	w.WriteHeader(http.StatusInternalServerError)
	w.Write([]byte(e.Error()))
	log.Println("error:", e.Error())
}

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(64)
	}

	storagepath := flag.Arg(0)

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		filepath := storagepath + r.RequestURI
		switch r.Method {
		case "GET":
			file, err := os.Open(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			_, err = io.Copy(w, file)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("GET", r.RequestURI)
		case "PUT":
			os.MkdirAll(path.Dir(filepath), 0755)
			file, err := os.Create(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			_, err = io.Copy(file, r.Body)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("PUT", r.RequestURI)
		case "DELETE":
			err := os.RemoveAll(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("DELETE", r.RequestURI)
		default:
			w.WriteHeader(http.StatusMethodNotAllowed)
		}
	})
	log.Println("Shelf serving files on " + *port + " from " + storagepath)
	log.Fatal(http.ListenAndServe(":"+*port, nil))
}
<commit_msg>providing default storage path, making it an optional flag. passing it as argument still works but is deprecated<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
)

var port = flag.String("p", "8888", "Port to listen on")
var storage = flag.String("s", "/var/lib/shelf", "Path to store files")

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage:	%v -p <port> -s <storage-path>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func errorResponse(w http.ResponseWriter, e error) {
	w.WriteHeader(http.StatusInternalServerError)
	w.Write([]byte(e.Error()))
	log.Println("error:", e.Error())
}

func main() {
	flag.Parse()

	var storagepath string
	if flag.NArg() == 1 {
		// deprecated: passing storage path as argument
		storagepath = flag.Arg(0)
	} else {
		storagepath = *storage
	}
	os.MkdirAll(storagepath, 0755)

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		filepath := storagepath + r.RequestURI
		switch r.Method {
		case "GET":
			file, err := os.Open(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			_, err = io.Copy(w, file)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("GET", r.RequestURI)
		case "PUT":
			os.MkdirAll(path.Dir(filepath), 0755)
			file, err := os.Create(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			_, err = io.Copy(file, r.Body)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("PUT", r.RequestURI)
		case "DELETE":
			err := os.RemoveAll(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("DELETE", r.RequestURI)
		default:
			w.WriteHeader(http.StatusMethodNotAllowed)
		}
	})
	log.Println("Shelf serving files on " + *port + " from " + storagepath)
	log.Fatal(http.ListenAndServe(":"+*port, nil))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/peterh/liner"
)

type Shell struct {
	shell  *liner.State
	prompt string
	cmds   map[string]shellCmd
	hist   string
	motor  Motor
}

func NewShell() *Shell {
	sh := &Shell{
		shell:  liner.NewLiner(),
		prompt: "mbus> ",
		hist:   filepath.Join(".", ".fcs_lpc_motor_history"),
		motor:  NewMotor("134.158.125.223:502"),
	}

	sh.shell.SetCtrlCAborts(true)
	sh.shell.SetCompleter(func(line string) (c []string) {
		for n := range sh.cmds {
			if strings.HasPrefix(n, strings.ToLower(line)) {
				c = append(c, n)
			}
		}
		return
	})

	if f, err := os.Open(sh.hist); err == nil {
		sh.shell.ReadHistory(f)
		f.Close()
	}

	sh.cmds = map[string]shellCmd{
		"dump":  sh.cmdDump,
		"get":   sh.cmdGet,
		"motor": sh.cmdMotor,
		"quit":  sh.cmdQuit,
		"set":   sh.cmdSet,
	}
	return sh
}

type shellCmd func(args []string) error

func (sh *Shell) Close() error {
	if f, err := os.Create(sh.hist); err != nil {
		log.Print("error writing history file: ", err)
	} else {
		sh.shell.WriteHistory(f)
		f.Close()
	}
	fmt.Printf("\n")
	return sh.shell.Close()
}

func (sh *Shell) run() error {
	for {
		raw, err := sh.shell.Prompt(sh.prompt)
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return err
		}
		// log.Printf("got: %q\n", raw)
		raw = strings.TrimSpace(raw)
		if raw == "" {
			continue
		}
		toks := strings.Split(raw, " ")
		err = sh.dispatch(toks)
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return err
		}
		sh.shell.AppendHistory(raw)
	}

	return nil
}

func (sh *Shell) dispatch(toks []string) error {
	var err error
	fct, ok := sh.cmds[toks[0]]
	if !ok {
		err = fmt.Errorf("invalid command [%s]", toks[0])
		log.Printf("error: %v\n", err)
		return err
	}

	return fct(toks[1:])
}

func (sh *Shell) cmdQuit(args []string) error {
	return io.EOF
}

func (sh *Shell) cmdGet(args []string) error {
	param, err := sh.parseParam(args[0])
	if err != nil {
		return err
	}

	o, err := sh.motor.read(param)
	if err != nil {
		return err
	}

	log.Printf(
		"Pr-%v: %s (%v)\n",
		param,
		displayBytes(o),
		codec.Uint16(o),
	)

	return err
}

func (sh *Shell) cmdSet(args []string) error {
	log.Printf(">>> %v\n", args)
	param, err := sh.parseParam(args[0])
	if err != nil {
		return err
	}
	vtype := "u16"
	if len(args) > 2 {
		vtype = args[2]
	}

	v := make([]byte, 2)

	switch vtype {
	case "u16", "uint16":
		vv, err := strconv.ParseUint(args[1], 10, 16)
		if err != nil {
			return err
		}
		codec.PutUint16(v, uint16(vv))

	default:
		return fmt.Errorf("cmd-set: invalid value-type (%v)", vtype)
	}

	log.Printf("set Pr-%v %s (%v)...\n", param, args[1], displayBytes(v))
	o, err := sh.motor.write(param, v)
	if err != nil {
		return err
	}
	log.Printf(
		"Pr-%v: %s (%v)\n",
		param,
		displayBytes(o),
		codec.Uint16(o),
	)

	return err
}

func (sh *Shell) cmdDump(args []string) error {
	var err error
	return err
}

func (sh *Shell) cmdMotor(args []string) error {
	switch len(args) {
	case 0:
		log.Printf("connected to [%s]\n", sh.motor.Address)
		return nil
	case 1:
		sh.motor = NewMotor(args[0])
		return nil
	default:
		return fmt.Errorf("cmd-motor: too many arguments (%d)", len(args))
	}
	return nil
}

func (sh *Shell) parseParam(arg string) (Parameter, error) {
	var err error
	var p Parameter

	if strings.Contains(arg, ".") {
		return NewParameterFromMenu(arg)
	}

	var reg uint64
	var base = 10
	if strings.HasPrefix(arg, "0x") {
		base = 16
		arg = arg[len("0x"):]
	}
	reg, err = strconv.ParseUint(arg, base, 64)
	if err != nil {
		return p, err
	}
	p = NewParameter(uint16(reg))
	return p, err
}

func displayBytes(o []byte) string {
	hex := make([]string, len(o))
	dec := make([]string, len(o))
	for i, v := range o {
		hex[i] = fmt.Sprintf("0x%02x", v)
		dec[i] = fmt.Sprintf("%3d", v)
	}

	return fmt.Sprintf("hex=%s dec=%s", hex, dec)
}
<commit_msg>shell: do not exit shell if get-cmd failed<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/peterh/liner"
)

type Shell struct {
	shell  *liner.State
	prompt string
	cmds   map[string]shellCmd
	hist   string
	motor  Motor
}

func NewShell() *Shell {
	sh := &Shell{
		shell:  liner.NewLiner(),
		prompt: "mbus> ",
		hist:   filepath.Join(".", ".fcs_lpc_motor_history"),
		motor:  NewMotor("134.158.125.223:502"),
	}

	sh.shell.SetCtrlCAborts(true)
	sh.shell.SetCompleter(func(line string) (c []string) {
		for n := range sh.cmds {
			if strings.HasPrefix(n, strings.ToLower(line)) {
				c = append(c, n)
			}
		}
		return
	})

	if f, err := os.Open(sh.hist); err == nil {
		sh.shell.ReadHistory(f)
		f.Close()
	}

	sh.cmds = map[string]shellCmd{
		"dump":  sh.cmdDump,
		"get":   sh.cmdGet,
		"motor": sh.cmdMotor,
		"quit":  sh.cmdQuit,
		"set":   sh.cmdSet,
	}
	return sh
}

type shellCmd func(args []string) error

func (sh *Shell) Close() error {
	if f, err := os.Create(sh.hist); err != nil {
		log.Print("error writing history file: ", err)
	} else {
		sh.shell.WriteHistory(f)
		f.Close()
	}
	fmt.Printf("\n")
	return sh.shell.Close()
}

func (sh *Shell) run() error {
	for {
		raw, err := sh.shell.Prompt(sh.prompt)
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return err
		}
		// log.Printf("got: %q\n", raw)
		raw = strings.TrimSpace(raw)
		if raw == "" {
			continue
		}
		toks := strings.Split(raw, " ")
		err = sh.dispatch(toks)
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return err
		}
		sh.shell.AppendHistory(raw)
	}

	return nil
}

func (sh *Shell) dispatch(toks []string) error {
	var err error
	fct, ok := sh.cmds[toks[0]]
	if !ok {
		err = fmt.Errorf("invalid command [%s]", toks[0])
		log.Printf("error: %v\n", err)
		return err
	}

	return fct(toks[1:])
}

func (sh *Shell) cmdQuit(args []string) error {
	return io.EOF
}

func (sh *Shell) cmdGet(args []string) error {
	param, err := sh.parseParam(args[0])
	if err != nil {
		return err
	}

	o, err := sh.motor.read(param)
	if err != nil {
		log.Printf("error: %v\n", err)
		err = nil
		return err
	}

	log.Printf(
		"Pr-%v: %s (%v)\n",
		param,
		displayBytes(o),
		codec.Uint16(o),
	)

	return err
}

func (sh *Shell) cmdSet(args []string) error {
	log.Printf(">>> %v\n", args)
	param, err := sh.parseParam(args[0])
	if err != nil {
		return err
	}
	vtype := "u16"
	if len(args) > 2 {
		vtype = args[2]
	}

	v := make([]byte, 2)

	switch vtype {
	case "u16", "uint16":
		vv, err := strconv.ParseUint(args[1], 10, 16)
		if err != nil {
			return err
		}
		codec.PutUint16(v, uint16(vv))

	default:
		return fmt.Errorf("cmd-set: invalid value-type (%v)", vtype)
	}

	log.Printf("set Pr-%v %s (%v)...\n", param, args[1], displayBytes(v))
	o, err := sh.motor.write(param, v)
	if err != nil {
		return err
	}
	log.Printf(
		"Pr-%v: %s (%v)\n",
		param,
		displayBytes(o),
		codec.Uint16(o),
	)

	return err
}

func (sh *Shell) cmdDump(args []string) error {
	var err error
	return err
}

func (sh *Shell) cmdMotor(args []string) error {
	switch len(args) {
	case 0:
		log.Printf("connected to [%s]\n", sh.motor.Address)
		return nil
	case 1:
		sh.motor = NewMotor(args[0])
		return nil
	default:
		return fmt.Errorf("cmd-motor: too many arguments (%d)", len(args))
	}
	return nil
}

func (sh *Shell) parseParam(arg string) (Parameter, error) {
	var err error
	var p Parameter

	if strings.Contains(arg, ".") {
		return NewParameterFromMenu(arg)
	}

	var reg uint64
	var base = 10
	if strings.HasPrefix(arg, "0x") {
		base = 16
		arg = arg[len("0x"):]
	}
	reg, err = strconv.ParseUint(arg, base, 64)
	if err != nil {
		return p, err
	}
	p = NewParameter(uint16(reg))
	return p, err
}

func displayBytes(o []byte) string {
	hex := make([]string, len(o))
	dec := make([]string, len(o))
	for i, v := range o {
		hex[i] = fmt.Sprintf("0x%02x", v)
		dec[i] = fmt.Sprintf("%3d", v)
	}

	return fmt.Sprintf("hex=%s dec=%s", hex, dec)
}
<|endoftext|>


<commit_before>package oauthutil

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"time"

	"github.com/ncw/rclone/fs"
	"github.com/skratchdot/open-golang/open"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

const (
	// ConfigToken is the key used to store the token under
	ConfigToken = "token"

	// ConfigClientID is the config key used to store the client id
	ConfigClientID = "client_id"

	// ConfigClientSecret is the config key used to store the client secret
	ConfigClientSecret = "client_secret"

	// TitleBarRedirectURL is the OAuth2 redirect URL to use when the authorization
	// code should be returned in the title bar of the browser, with the page text
	// prompting the user to copy the code and paste it in the application.
	TitleBarRedirectURL = "urn:ietf:wg:oauth:2.0:oob"

	// bindPort is the port that we bind the local webserver to
	bindPort = "53682"

	// bindAddress is binding for local webserver when active
	bindAddress = "127.0.0.1:" + bindPort

	// RedirectURL is redirect to local webserver when active
	RedirectURL = "http://" + bindAddress + "/"

	// RedirectPublicURL is redirect to local webserver when active with public name
	RedirectPublicURL = "http://localhost.rclone.org:" + bindPort + "/"
)

// oldToken contains an end-user's tokens.
// This is the data you must store to persist authentication.
//
// From the original code.google.com/p/goauth2/oauth package - used
// for backwards compatibility in the rclone config file
type oldToken struct {
	AccessToken  string
	RefreshToken string
	Expiry       time.Time
}

// getToken returns the token saved in the config file under
// section name.
func getToken(name string) (*oauth2.Token, error) {
	tokenString, err := fs.ConfigFile.GetValue(string(name), ConfigToken)
	if err != nil {
		return nil, err
	}
	if tokenString == "" {
		return nil, fmt.Errorf("Empty token found - please run rclone config again")
	}
	token := new(oauth2.Token)
	err = json.Unmarshal([]byte(tokenString), token)
	if err != nil {
		return nil, err
	}
	// if has data then return it
	if token.AccessToken != "" && token.RefreshToken != "" {
		return token, nil
	}
	// otherwise try parsing as oldToken
	oldtoken := new(oldToken)
	err = json.Unmarshal([]byte(tokenString), oldtoken)
	if err != nil {
		return nil, err
	}
	// Fill in result into new token
	token.AccessToken = oldtoken.AccessToken
	token.RefreshToken = oldtoken.RefreshToken
	token.Expiry = oldtoken.Expiry
	// Save new format in config file
	err = putToken(name, token)
	if err != nil {
		return nil, err
	}
	return token, nil
}

// putToken stores the token in the config file
//
// This saves the config file if it changes
func putToken(name string, token *oauth2.Token) error {
	tokenBytes, err := json.Marshal(token)
	if err != nil {
		return err
	}
	tokenString := string(tokenBytes)
	old := fs.ConfigFile.MustValue(name, ConfigToken)
	if tokenString != old {
		fs.ConfigFile.SetValue(name, ConfigToken, tokenString)
		fs.SaveConfig()
		fs.Debug(name, "Saving new token in config file")
	}
	return nil
}

// tokenSource stores updated tokens in the config file
type tokenSource struct {
	Name        string
	TokenSource oauth2.TokenSource
	OldToken    oauth2.Token
}

// Token returns a token or an error.
// Token must be safe for concurrent use by multiple goroutines.
// The returned Token must not be modified.
//
// This saves the token in the config file if it has changed
func (ts *tokenSource) Token() (*oauth2.Token, error) {
	token, err := ts.TokenSource.Token()
	if err != nil {
		return nil, err
	}
	if *token != ts.OldToken {
		err = putToken(ts.Name, token)
		if err != nil {
			return nil, err
		}
	}
	return token, nil
}

// Check interface satisfied
var _ oauth2.TokenSource = (*tokenSource)(nil)

// Context returns a context with our HTTP Client baked in for oauth2
func Context() context.Context {
	return context.WithValue(nil, oauth2.HTTPClient, fs.Config.Client())
}

// overrideCredentials sets the ClientID and ClientSecret from the
// config file if they are not blank
func overrideCredentials(name string, config *oauth2.Config) {
	ClientID := fs.ConfigFile.MustValue(name, ConfigClientID)
	if ClientID != "" {
		config.ClientID = ClientID
	}
	ClientSecret := fs.ConfigFile.MustValue(name, ConfigClientSecret)
	if ClientSecret != "" {
		config.ClientSecret = ClientSecret
	}
}

// NewClient gets a token from the config file and configures a Client
// with it
func NewClient(name string, config *oauth2.Config) (*http.Client, error) {
	overrideCredentials(name, config)
	token, err := getToken(name)
	if err != nil {
		return nil, err
	}

	// Set our own http client in the context
	ctx := Context()

	// Wrap the TokenSource in our TokenSource which saves changed
	// tokens in the config file
	ts := &tokenSource{
		Name:        name,
		OldToken:    *token,
		TokenSource: config.TokenSource(ctx, token),
	}
	return oauth2.NewClient(ctx, ts), nil

}

// Config does the initial creation of the token
//
// It may run an internal webserver to receive the results
func Config(name string, config *oauth2.Config) error {
	overrideCredentials(name, config)
	// See if already have a token
	tokenString := fs.ConfigFile.MustValue(name, "token")
	if tokenString != "" {
		fmt.Printf("Already have a token - refresh?\n")
		if !fs.Confirm() {
			return nil
		}
	}

	// Detect whether we should use internal web server
	useWebServer := false
	switch config.RedirectURL {
	case RedirectURL, RedirectPublicURL:
		useWebServer = true
	case TitleBarRedirectURL:
		fmt.Printf("Use auto config?\n")
		fmt.Printf(" * Say Y if not sure\n")
		fmt.Printf(" * Say N if you are working on a remote or headless machine or Y didn't work\n")
		useWebServer = fs.Confirm()
		if useWebServer {
			// copy the config and set to use the internal webserver
			configCopy := *config
			config = &configCopy
			config.RedirectURL = RedirectURL
		}
	}

	// Make random state
	stateBytes := make([]byte, 16)
	_, err := rand.Read(stateBytes)
	if err != nil {
		return err
	}
	state := fmt.Sprintf("%x", stateBytes)
	authURL := config.AuthCodeURL(state)

	// Prepare webserver
	server := authServer{
		state:       state,
		bindAddress: bindAddress,
		authURL:     authURL,
	}
	if useWebServer {
		server.code = make(chan string, 1)
		go server.Start()
		defer server.Stop()
		authURL = "http://" + bindAddress + "/auth"
	}

	// Generate a URL for the user to visit for authorization.
	_ = open.Start(authURL)
	fmt.Printf("If your browser doesn't open automatically go to the following link: %s\n", authURL)
	fmt.Printf("Log in and authorize rclone for access\n")

	var authCode string
	if useWebServer {
		// Read the code, and exchange it for a token.
		fmt.Printf("Waiting for code...\n")
		authCode = <-server.code
		if authCode != "" {
			fmt.Printf("Got code\n")
		} else {
			return fmt.Errorf("Failed to get code")
		}
	} else {
		// Read the code, and exchange it for a token.
		fmt.Printf("Enter verification code> ")
		authCode = fs.ReadLine()
	}
	token, err := config.Exchange(oauth2.NoContext, authCode)
	if err != nil {
		return fmt.Errorf("Failed to get token: %v", err)
	}
	return putToken(name, token)
}

// Local web server for collecting auth
type authServer struct {
	state       string
	listener    net.Listener
	bindAddress string
	code        chan string
	authURL     string
}

// startWebServer runs an internal web server to receive config details
func (s *authServer) Start() {
	fs.Debug(nil, "Starting auth server on %s", s.bindAddress)
	mux := http.NewServeMux()
	server := &http.Server{
		Addr:    s.bindAddress,
		Handler: mux,
	}
	mux.HandleFunc("/favicon.ico", func(w http.ResponseWriter, req *http.Request) {
		http.Error(w, "", 404)
		return
	})
	mux.HandleFunc("/auth", func(w http.ResponseWriter, req *http.Request) {
		http.Redirect(w, req, s.authURL, 307)
		return
	})
	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		fs.Debug(nil, "Received request on auth server")
		code := req.FormValue("code")
		if code != "" {
			state := req.FormValue("state")
			if state != s.state {
				fs.Debug(nil, "State did not match: want %q got %q", s.state, state)
				fmt.Fprintf(w, "
Failure
\n
Auth state doesn't match
")
			} else {
				fs.Debug(nil, "Successfully got code")
				if s.code != nil {
					fmt.Fprintf(w, "
Success
\n
Go back to rclone to continue
")
					s.code <- code
				} else {
					fmt.Fprintf(w, "
Success
\n
Cut and paste this code into rclone: %s
", code)
				}
			}
			return
		}
		fs.Debug(nil, "No code found on request")
		fmt.Fprintf(w, "
Failed!
\nNo code found.")
		http.Error(w, "", 500)
	})

	var err error
	s.listener, err = net.Listen("tcp", s.bindAddress)
	if err != nil {
		log.Fatalf("Failed to start auth webserver: %v", err)
	}
	err = server.Serve(s.listener)
	fs.Debug(nil, "Closed auth server with error: %v", err)
}

func (s *authServer) Stop() {
	fs.Debug(nil, "Closing auth server")
	if s.code != nil {
		close(s.code)
		s.code = nil
	}
	_ = s.listener.Close()
}
<commit_msg>Add RedirectLocalhostURL for another form of redirect URL<commit_after>package oauthutil

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"time"

	"github.com/ncw/rclone/fs"
	"github.com/skratchdot/open-golang/open"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

const (
	// ConfigToken is the key used to store the token under
	ConfigToken = "token"

	// ConfigClientID is the config key used to store the client id
	ConfigClientID = "client_id"

	// ConfigClientSecret is the config key used to store the client secret
	ConfigClientSecret = "client_secret"

	// TitleBarRedirectURL is the OAuth2 redirect URL to use when the authorization
	// code should be returned in the title bar of the browser, with the page text
	// prompting the user to copy the code and paste it in the application.
	TitleBarRedirectURL = "urn:ietf:wg:oauth:2.0:oob"

	// bindPort is the port that we bind the local webserver to
	bindPort = "53682"

	// bindAddress is binding for local webserver when active
	bindAddress = "127.0.0.1:" + bindPort

	// RedirectURL is redirect to local webserver when active
	RedirectURL = "http://" + bindAddress + "/"

	// RedirectPublicURL is redirect to local webserver when active with public name
	RedirectPublicURL = "http://localhost.rclone.org:" + bindPort + "/"

	// RedirectLocalhostURL is redirect to local webserver when active with localhost
	RedirectLocalhostURL = "http://localhost:" + bindPort + "/"
)

// oldToken contains an end-user's tokens.
// This is the data you must store to persist authentication.
//
// From the original code.google.com/p/goauth2/oauth package - used
// for backwards compatibility in the rclone config file
type oldToken struct {
	AccessToken  string
	RefreshToken string
	Expiry       time.Time
}

// getToken returns the token saved in the config file under
// section name.
func getToken(name string) (*oauth2.Token, error) {
	tokenString, err := fs.ConfigFile.GetValue(string(name), ConfigToken)
	if err != nil {
		return nil, err
	}
	if tokenString == "" {
		return nil, fmt.Errorf("Empty token found - please run rclone config again")
	}
	token := new(oauth2.Token)
	err = json.Unmarshal([]byte(tokenString), token)
	if err != nil {
		return nil, err
	}
	// if has data then return it
	if token.AccessToken != "" && token.RefreshToken != "" {
		return token, nil
	}
	// otherwise try parsing as oldToken
	oldtoken := new(oldToken)
	err = json.Unmarshal([]byte(tokenString), oldtoken)
	if err != nil {
		return nil, err
	}
	// Fill in result into new token
	token.AccessToken = oldtoken.AccessToken
	token.RefreshToken = oldtoken.RefreshToken
	token.Expiry = oldtoken.Expiry
	// Save new format in config file
	err = putToken(name, token)
	if err != nil {
		return nil, err
	}
	return token, nil
}

// putToken stores the token in the config file
//
// This saves the config file if it changes
func putToken(name string, token *oauth2.Token) error {
	tokenBytes, err := json.Marshal(token)
	if err != nil {
		return err
	}
	tokenString := string(tokenBytes)
	old := fs.ConfigFile.MustValue(name, ConfigToken)
	if tokenString != old {
		fs.ConfigFile.SetValue(name, ConfigToken, tokenString)
		fs.SaveConfig()
		fs.Debug(name, "Saving new token in config file")
	}
	return nil
}

// tokenSource stores updated tokens in the config file
type tokenSource struct {
	Name        string
	TokenSource oauth2.TokenSource
	OldToken    oauth2.Token
}

// Token returns a token or an error.
// Token must be safe for concurrent use by multiple goroutines.
// The returned Token must not be modified.
//
// This saves the token in the config file if it has changed
func (ts *tokenSource) Token() (*oauth2.Token, error) {
	token, err := ts.TokenSource.Token()
	if err != nil {
		return nil, err
	}
	if *token != ts.OldToken {
		err = putToken(ts.Name, token)
		if err != nil {
			return nil, err
		}
	}
	return token, nil
}

// Check interface satisfied
var _ oauth2.TokenSource = (*tokenSource)(nil)

// Context returns a context with our HTTP Client baked in for oauth2
func Context() context.Context {
	return context.WithValue(nil, oauth2.HTTPClient, fs.Config.Client())
}

// overrideCredentials sets the ClientID and ClientSecret from the
// config file if they are not blank
func overrideCredentials(name string, config *oauth2.Config) {
	ClientID := fs.ConfigFile.MustValue(name, ConfigClientID)
	if ClientID != "" {
		config.ClientID = ClientID
	}
	ClientSecret := fs.ConfigFile.MustValue(name, ConfigClientSecret)
	if ClientSecret != "" {
		config.ClientSecret = ClientSecret
	}
}

// NewClient gets a token from the config file and configures a Client
// with it
func NewClient(name string, config *oauth2.Config) (*http.Client, error) {
	overrideCredentials(name, config)
	token, err := getToken(name)
	if err != nil {
		return nil, err
	}

	// Set our own http client in the context
	ctx := Context()

	// Wrap the TokenSource in our TokenSource which saves changed
	// tokens in the config file
	ts := &tokenSource{
		Name:        name,
		OldToken:    *token,
		TokenSource: config.TokenSource(ctx, token),
	}
	return oauth2.NewClient(ctx, ts), nil

}

// Config does the initial creation of the token
//
// It may run an internal webserver to receive the results
func Config(name string, config *oauth2.Config) error {
	overrideCredentials(name, config)
	// See if already have a token
	tokenString := fs.ConfigFile.MustValue(name, "token")
	if tokenString != "" {
		fmt.Printf("Already have a token - refresh?\n")
		if !fs.Confirm() {
			return nil
		}
	}

	// Detect whether we should use internal web server
	useWebServer := false
	switch config.RedirectURL {
	case RedirectURL, RedirectPublicURL, RedirectLocalhostURL:
		useWebServer = true
	case TitleBarRedirectURL:
		fmt.Printf("Use auto config?\n")
		fmt.Printf(" * Say Y if not sure\n")
		fmt.Printf(" * Say N if you are working on a remote or headless machine or Y didn't work\n")
		useWebServer = fs.Confirm()
		if useWebServer {
			// copy the config and set to use the internal webserver
			configCopy := *config
			config = &configCopy
			config.RedirectURL = RedirectURL
		}
	}

	// Make random state
	stateBytes := make([]byte, 16)
	_, err := rand.Read(stateBytes)
	if err != nil {
		return err
	}
	state := fmt.Sprintf("%x", stateBytes)
	authURL := config.AuthCodeURL(state)

	// Prepare webserver
	server := authServer{
		state:       state,
		bindAddress: bindAddress,
		authURL:     authURL,
	}
	if useWebServer {
		server.code = make(chan string, 1)
		go server.Start()
		defer server.Stop()
		authURL = "http://" + bindAddress + "/auth"
	}

	// Generate a URL for the user to visit for authorization.
	_ = open.Start(authURL)
	fmt.Printf("If your browser doesn't open automatically go to the following link: %s\n", authURL)
	fmt.Printf("Log in and authorize rclone for access\n")

	var authCode string
	if useWebServer {
		// Read the code, and exchange it for a token.
		fmt.Printf("Waiting for code...\n")
		authCode = <-server.code
		if authCode != "" {
			fmt.Printf("Got code\n")
		} else {
			return fmt.Errorf("Failed to get code")
		}
	} else {
		// Read the code, and exchange it for a token.
		fmt.Printf("Enter verification code> ")
		authCode = fs.ReadLine()
	}
	token, err := config.Exchange(oauth2.NoContext, authCode)
	if err != nil {
		return fmt.Errorf("Failed to get token: %v", err)
	}
	return putToken(name, token)
}

// Local web server for collecting auth
type authServer struct {
	state       string
	listener    net.Listener
	bindAddress string
	code        chan string
	authURL     string
}

// startWebServer runs an internal web server to receive config details
func (s *authServer) Start() {
	fs.Debug(nil, "Starting auth server on %s", s.bindAddress)
	mux := http.NewServeMux()
	server := &http.Server{
		Addr:    s.bindAddress,
		Handler: mux,
	}
	mux.HandleFunc("/favicon.ico", func(w http.ResponseWriter, req *http.Request) {
		http.Error(w, "", 404)
		return
	})
	mux.HandleFunc("/auth", func(w http.ResponseWriter, req *http.Request) {
		http.Redirect(w, req, s.authURL, 307)
		return
	})
	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		fs.Debug(nil, "Received request on auth server")
		code := req.FormValue("code")
		if code != "" {
			state := req.FormValue("state")
			if state != s.state {
				fs.Debug(nil, "State did not match: want %q got %q", s.state, state)
				fmt.Fprintf(w, "
Failure
\n
Auth state doesn't match
")
			} else {
				fs.Debug(nil, "Successfully got code")
				if s.code != nil {
					fmt.Fprintf(w, "
Success
\n
Go back to rclone to continue
")
					s.code <- code
				} else {
					fmt.Fprintf(w, "
Success
\n
Cut and paste this code into rclone: %s
", code)
				}
			}
			return
		}
		fs.Debug(nil, "No code found on request")
		fmt.Fprintf(w, "
Failed!
\nNo code found.")
		http.Error(w, "", 500)
	})

	var err error
	s.listener, err = net.Listen("tcp", s.bindAddress)
	if err != nil {
		log.Fatalf("Failed to start auth webserver: %v", err)
	}
	err = server.Serve(s.listener)
	fs.Debug(nil, "Closed auth server with error: %v", err)
}

func (s *authServer) Stop() {
	fs.Debug(nil, "Closing auth server")
	if s.code != nil {
		close(s.code)
		s.code = nil
	}
	_ = s.listener.Close()
}
<|endoftext|>


<commit_before>// package shell implements a remote API interface for a running ipfs daemon
package shell

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	gohttp "net/http"
	"os"

	files "github.com/whyrusleeping/go-multipart-files"
	tar "github.com/whyrusleeping/tar-utils"
)

type Shell struct {
	url     string
	httpcli *gohttp.Client
}

func NewShell(url string) *Shell {
	return &Shell{
		url: url,
		httpcli: &gohttp.Client{
			Transport: &gohttp.Transport{
				DisableKeepAlives: true,
			},
		},
	}
}

func (s *Shell) newRequest(command string, args ...string) *Request {
	return NewRequest(s.url, command, args...)
}

type IdOutput struct {
	ID              string
	PublicKey       string
	Addresses       []string
	AgentVersion    string
	ProtocolVersion string
}

// ID gets information about a given peer.  Arguments:
//
// peer: peer.ID of the node to look up.  If no peer is specified,
//   return information about the local peer.
func (s *Shell) ID(peer ...string) (*IdOutput, error) {
	if len(peer) > 1 {
		return nil, fmt.Errorf("Too many peer arguments")
	}

	resp, err := NewRequest(s.url, "id", peer...).Send(s.httpcli)
	if err != nil {
		return nil, err
	}

	defer resp.Close()
	if resp.Error != nil {
		return nil, resp.Error
	}

	decoder := json.NewDecoder(resp.Output)
	out := new(IdOutput)
	err = decoder.Decode(out)
	if err != nil {
		return nil, err
	}

	return out, nil
}

// Cat the content at the given path. Callers need to drain and close the returned reader after usage.
func (s *Shell) Cat(path string) (io.ReadCloser, error) {
	resp, err := NewRequest(s.url, "cat", path).Send(s.httpcli)
	if err != nil {
		return nil, err
	}
	if resp.Error != nil {
		return nil, resp.Error
	}

	return resp.Output, nil
}

type object struct {
	Hash string
}

// Add a file to ipfs from the given reader, returns the hash of the added file
func (s *Shell) Add(r io.Reader) (string, error) {
	var rc io.ReadCloser
	if rclose, ok := r.(io.ReadCloser); ok {
		rc = rclose
	} else {
		rc = ioutil.NopCloser(r)
	}

	// handler expects an array of files
	fr := files.NewReaderFile("", "", rc, nil)
	slf := files.NewSliceFile("", "", []files.File{fr})
	fileReader := files.NewMultiFileReader(slf, true)

	req := NewRequest(s.url, "add")
	req.Body = fileReader

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()
	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

func (s *Shell) AddLink(target string) (string, error) {
	link := files.NewLinkFile("", "", target, nil)
	slf := files.NewSliceFile("", "", []files.File{link})
	reader := files.NewMultiFileReader(slf, true)

	req := s.newRequest("add")
	req.Body = reader

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()
	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

// AddDir adds a directory recursively with all of the files under it
func (s *Shell) AddDir(dir string) (string, error) {
	stat, err := os.Lstat(dir)
	if err != nil {
		return "", err
	}

	sf, err := files.NewSerialFile("", dir, stat)
	if err != nil {
		return "", err
	}
	slf := files.NewSliceFile("", dir, []files.File{sf})
	reader := files.NewMultiFileReader(slf, true)

	req := NewRequest(s.url, "add")
	req.Opts["r"] = "true"
	req.Body = reader

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()
	if resp.Error != nil {
		return "", resp.Error
	}

	dec := json.NewDecoder(resp.Output)
	var final string
	for {
		var out object
		err = dec.Decode(&out)
		if err != nil {
			if err == io.EOF {
				break
			}
			return "", err
		}
		final = out.Hash
	}

	if final == "" {
		return "", errors.New("no results received")
	}

	return final, nil
}

const (
	TRaw = iota
	TDirectory
	TFile
	TMetadata
	TSymlink
)

// List entries at the given path
func (s *Shell) List(path string) ([]*LsLink, error) {
	resp, err := NewRequest(s.url, "ls", path).Send(s.httpcli)
	if err != nil {
		return nil, err
	}
	defer resp.Close()

	if resp.Error != nil {
		return nil, resp.Error
	}

	var out struct{ Objects []LsObject }
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return nil, err
	}

	return out.Objects[0].Links, nil
}

type LsLink struct {
	Hash string
	Name string
	Size uint64
	Type int
}

type LsObject struct {
	Links []*LsLink
	LsLink
}

// Pin the given path
func (s *Shell) Pin(path string) error {
	req := NewRequest(s.url, "pin/add", path)
	req.Opts["r"] = "true"

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return err
	}
	defer resp.Close()
	if resp.Error != nil {
		return resp.Error
	}

	return nil
}

// Unpin the given path
func (s *Shell) Unpin(path string) error {
	req := NewRequest(s.url, "pin/rm", path)
	req.Opts["r"] = "true"

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return err
	}
	defer resp.Close()
	if resp.Error != nil {
		return resp.Error
	}

	return nil
}

type PeerInfo struct {
	Addrs []string
	ID    string
}

func (s *Shell) FindPeer(peer string) (*PeerInfo, error) {
	resp, err := s.newRequest("dht/findpeer", peer).Send(s.httpcli)
	if err != nil {
		return nil, err
	}
	defer resp.Close()

	if resp.Error != nil {
		return nil, resp.Error
	}

	str := struct{ Responses []PeerInfo }{}
	err = json.NewDecoder(resp.Output).Decode(&str)
	if err != nil {
		return nil, err
	}

	if len(str.Responses) == 0 {
		return nil, errors.New("peer not found")
	}

	return &str.Responses[0], nil
}

func (s *Shell) Refs(hash string, recursive bool) (<-chan string, error) {
	req := s.newRequest("refs", hash)
	if recursive {
		req.Opts["r"] = "true"
	}

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return nil, err
	}

	if resp.Error != nil {
		return nil, resp.Error
	}

	out := make(chan string)
	go func() {
		defer resp.Close()
		scan := bufio.NewScanner(resp.Output)
		for scan.Scan() {
			if len(scan.Text()) > 0 {
				out <- scan.Text()
			}
		}
		close(out)
	}()

	return out, nil
}

func (s *Shell) Patch(root, action string, args ...string) (string, error) {
	cmdargs := append([]string{root}, args...)
	resp, err := s.newRequest("object/patch/"+action, cmdargs...).Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", resp.Error
	}

	dec := json.NewDecoder(resp.Output)
	var out object
	err = dec.Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

func (s *Shell) PatchLink(root, path, childhash string, create bool) (string, error) {
	cmdargs := []string{root, path, childhash}

	req := s.newRequest("object/patch/add-link", cmdargs...)
	if create {
		req.Opts["create"] = "true"
	}

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

func (s *Shell) Get(hash, outdir string) error {
	resp, err := s.newRequest("get", hash).Send(s.httpcli)
	if err != nil {
		return err
	}
	defer resp.Close()

	if resp.Error != nil {
		return resp.Error
	}

	extractor := &tar.Extractor{Path: outdir}
	return extractor.Extract(resp.Output)
}

func (s *Shell) NewObject(template string) (string, error) {
	args := []string{}
	if template != "" {
		args = []string{template}
	}

	resp, err := s.newRequest("object/new", args...).Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

func (s *Shell) PutObject(encoding string, r io.Reader) (string, error) {
	var rc io.ReadCloser
	if rclose, ok := r.(io.ReadCloser); ok {
		rc = rclose
	} else {
		rc = ioutil.NopCloser(r)
	}

	// handler expects an array of files
	fr := files.NewReaderFile("", "", rc, nil)
	slf := files.NewSliceFile("", "", []files.File{fr})
	fileReader := files.NewMultiFileReader(slf, true)

	req := NewRequest(s.url, "object/put", encoding)
	req.Body = fileReader

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()
	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

func (s *Shell) ResolvePath(path string) (string, error) {
	resp, err := s.newRequest("object/stat", path).Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

// returns ipfs version and commit sha
func (s *Shell) Version() (string, string, error) {
	resp, err := s.newRequest("version").Send(s.httpcli)
	if err != nil {
		return "", "", err
	}

	defer resp.Close()
	if resp.Error != nil {
		return "", "", resp.Error
	}

	ver := struct {
		Version string
		Commit  string
	}{}

	err = json.NewDecoder(resp.Output).Decode(&ver)
	if err != nil {
		return "", "", err
	}

	return ver.Version, ver.Commit, nil
}

func (s *Shell) IsUp() bool {
	_, _, err := s.Version()
	return err == nil
}

func (s *Shell) BlockStat(path string) (string, int, error) {
	resp, err := s.newRequest("block/stat", path).Send(s.httpcli)
	if err != nil {
		return "", 0, err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", 0, resp.Error
	}

	var inf struct {
		Key  string
		Size int
	}

	err = json.NewDecoder(resp.Output).Decode(&inf)
	if err != nil {
		return "", 0, err
	}

	return inf.Key, inf.Size, nil
}
<commit_msg>added GetObject func<commit_after>// package shell implements a remote API interface for a running ipfs daemon
package shell

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	gohttp "net/http"
	"os"

	files "github.com/whyrusleeping/go-multipart-files"
	tar "github.com/whyrusleeping/tar-utils"
)

type Shell struct {
	url     string
	httpcli *gohttp.Client
}

func NewShell(url string) *Shell {
	return &Shell{
		url: url,
		httpcli: &gohttp.Client{
			Transport: &gohttp.Transport{
				DisableKeepAlives: true,
			},
		},
	}
}

func (s *Shell) newRequest(command string, args ...string) *Request {
	return NewRequest(s.url, command, args...)
}

type IdOutput struct {
	ID              string
	PublicKey       string
	Addresses       []string
	AgentVersion    string
	ProtocolVersion string
}

// ID gets information about a given peer.  Arguments:
//
// peer: peer.ID of the node to look up.  If no peer is specified,
//   return information about the local peer.
func (s *Shell) ID(peer ...string) (*IdOutput, error) {
	if len(peer) > 1 {
		return nil, fmt.Errorf("Too many peer arguments")
	}

	resp, err := NewRequest(s.url, "id", peer...).Send(s.httpcli)
	if err != nil {
		return nil, err
	}

	defer resp.Close()
	if resp.Error != nil {
		return nil, resp.Error
	}

	decoder := json.NewDecoder(resp.Output)
	out := new(IdOutput)
	err = decoder.Decode(out)
	if err != nil {
		return nil, err
	}

	return out, nil
}

// Cat the content at the given path. Callers need to drain and close the returned reader after usage.
func (s *Shell) Cat(path string) (io.ReadCloser, error) {
	resp, err := NewRequest(s.url, "cat", path).Send(s.httpcli)
	if err != nil {
		return nil, err
	}
	if resp.Error != nil {
		return nil, resp.Error
	}

	return resp.Output, nil
}

type object struct {
	Hash string
}

// Add a file to ipfs from the given reader, returns the hash of the added file
func (s *Shell) Add(r io.Reader) (string, error) {
	var rc io.ReadCloser
	if rclose, ok := r.(io.ReadCloser); ok {
		rc = rclose
	} else {
		rc = ioutil.NopCloser(r)
	}

	// handler expects an array of files
	fr := files.NewReaderFile("", "", rc, nil)
	slf := files.NewSliceFile("", "", []files.File{fr})
	fileReader := files.NewMultiFileReader(slf, true)

	req := NewRequest(s.url, "add")
	req.Body = fileReader

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()
	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

func (s *Shell) AddLink(target string) (string, error) {
	link := files.NewLinkFile("", "", target, nil)
	slf := files.NewSliceFile("", "", []files.File{link})
	reader := files.NewMultiFileReader(slf, true)

	req := s.newRequest("add")
	req.Body = reader

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()
	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

// AddDir adds a directory recursively with all of the files under it
func (s *Shell) AddDir(dir string) (string, error) {
	stat, err := os.Lstat(dir)
	if err != nil {
		return "", err
	}

	sf, err := files.NewSerialFile("", dir, stat)
	if err != nil {
		return "", err
	}
	slf := files.NewSliceFile("", dir, []files.File{sf})
	reader := files.NewMultiFileReader(slf, true)

	req := NewRequest(s.url, "add")
	req.Opts["r"] = "true"
	req.Body = reader

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()
	if resp.Error != nil {
		return "", resp.Error
	}

	dec := json.NewDecoder(resp.Output)
	var final string
	for {
		var out object
		err = dec.Decode(&out)
		if err != nil {
			if err == io.EOF {
				break
			}
			return "", err
		}
		final = out.Hash
	}

	if final == "" {
		return "", errors.New("no results received")
	}

	return final, nil
}

const (
	TRaw = iota
	TDirectory
	TFile
	TMetadata
	TSymlink
)

// List entries at the given path
func (s *Shell) List(path string) ([]*LsLink, error) {
	resp, err := NewRequest(s.url, "ls", path).Send(s.httpcli)
	if err != nil {
		return nil, err
	}
	defer resp.Close()

	if resp.Error != nil {
		return nil, resp.Error
	}

	var out struct{ Objects []LsObject }
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return nil, err
	}

	return out.Objects[0].Links, nil
}

type LsLink struct {
	Hash string
	Name string
	Size uint64
	Type int
}

type LsObject struct {
	Links []*LsLink
	LsLink
}

// Pin the given path
func (s *Shell) Pin(path string) error {
	req := NewRequest(s.url, "pin/add", path)
	req.Opts["r"] = "true"

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return err
	}
	defer resp.Close()
	if resp.Error != nil {
		return resp.Error
	}

	return nil
}

// Unpin the given path
func (s *Shell) Unpin(path string) error {
	req := NewRequest(s.url, "pin/rm", path)
	req.Opts["r"] = "true"

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return err
	}
	defer resp.Close()
	if resp.Error != nil {
		return resp.Error
	}

	return nil
}

type PeerInfo struct {
	Addrs []string
	ID    string
}

func (s *Shell) FindPeer(peer string) (*PeerInfo, error) {
	resp, err := s.newRequest("dht/findpeer", peer).Send(s.httpcli)
	if err != nil {
		return nil, err
	}
	defer resp.Close()

	if resp.Error != nil {
		return nil, resp.Error
	}

	str := struct{ Responses []PeerInfo }{}
	err = json.NewDecoder(resp.Output).Decode(&str)
	if err != nil {
		return nil, err
	}

	if len(str.Responses) == 0 {
		return nil, errors.New("peer not found")
	}

	return &str.Responses[0], nil
}

func (s *Shell) Refs(hash string, recursive bool) (<-chan string, error) {
	req := s.newRequest("refs", hash)
	if recursive {
		req.Opts["r"] = "true"
	}

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return nil, err
	}

	if resp.Error != nil {
		return nil, resp.Error
	}

	out := make(chan string)
	go func() {
		defer resp.Close()
		scan := bufio.NewScanner(resp.Output)
		for scan.Scan() {
			if len(scan.Text()) > 0 {
				out <- scan.Text()
			}
		}
		close(out)
	}()

	return out, nil
}

func (s *Shell) Patch(root, action string, args ...string) (string, error) {
	cmdargs := append([]string{root}, args...)
	resp, err := s.newRequest("object/patch/"+action, cmdargs...).Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", resp.Error
	}

	dec := json.NewDecoder(resp.Output)
	var out object
	err = dec.Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

func (s *Shell) PatchLink(root, path, childhash string, create bool) (string, error) {
	cmdargs := []string{root, path, childhash}

	req := s.newRequest("object/patch/add-link", cmdargs...)
	if create {
		req.Opts["create"] = "true"
	}

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

func (s *Shell) Get(hash, outdir string) error {
	resp, err := s.newRequest("get", hash).Send(s.httpcli)
	if err != nil {
		return err
	}
	defer resp.Close()

	if resp.Error != nil {
		return resp.Error
	}

	extractor := &tar.Extractor{Path: outdir}
	return extractor.Extract(resp.Output)
}

func (s *Shell) NewObject(template string) (string, error) {
	args := []string{}
	if template != "" {
		args = []string{template}
	}

	resp, err := s.newRequest("object/new", args...).Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

func (s *Shell) PutObject(encoding string, r io.Reader) (string, error) {
	var rc io.ReadCloser
	if rclose, ok := r.(io.ReadCloser); ok {
		rc = rclose
	} else {
		rc = ioutil.NopCloser(r)
	}

	// handler expects an array of files
	fr := files.NewReaderFile("", "", rc, nil)
	slf := files.NewSliceFile("", "", []files.File{fr})
	fileReader := files.NewMultiFileReader(slf, true)

	req := NewRequest(s.url, "object/put", encoding)
	req.Body = fileReader

	resp, err := req.Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()
	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

type GetObjectJson struct {
    Data string  `json:"data"`
    Links []Link `json:"links"`
}

type Link struct {
    Hash string `json:"hash"`
    Name string `json:"name"`
    Size uint64 `json:"size"`
}

func (s *Shell) GetObject(hash string) (string, error) {
	resp, err := s.newRequest("object/get", hash).Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", resp.Error
	}

	var out GetObjectJson
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}
	json, err := json.Marshal(out)
	if err != nil {
		return "", err
	}

	return string(json), nil
}

func (s *Shell) ResolvePath(path string) (string, error) {
	resp, err := s.newRequest("object/stat", path).Send(s.httpcli)
	if err != nil {
		return "", err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", resp.Error
	}

	var out object
	err = json.NewDecoder(resp.Output).Decode(&out)
	if err != nil {
		return "", err
	}

	return out.Hash, nil
}

// returns ipfs version and commit sha
func (s *Shell) Version() (string, string, error) {
	resp, err := s.newRequest("version").Send(s.httpcli)
	if err != nil {
		return "", "", err
	}

	defer resp.Close()
	if resp.Error != nil {
		return "", "", resp.Error
	}

	ver := struct {
		Version string
		Commit  string
	}{}

	err = json.NewDecoder(resp.Output).Decode(&ver)
	if err != nil {
		return "", "", err
	}

	return ver.Version, ver.Commit, nil
}

func (s *Shell) IsUp() bool {
	_, _, err := s.Version()
	return err == nil
}

func (s *Shell) BlockStat(path string) (string, int, error) {
	resp, err := s.newRequest("block/stat", path).Send(s.httpcli)
	if err != nil {
		return "", 0, err
	}
	defer resp.Close()

	if resp.Error != nil {
		return "", 0, resp.Error
	}

	var inf struct {
		Key  string
		Size int
	}

	err = json.NewDecoder(resp.Output).Decode(&inf)
	if err != nil {
		return "", 0, err
	}

	return inf.Key, inf.Size, nil
}
<|endoftext|>


<commit_before>// Session implements middleware for easily using github.com/gorilla/sessions
// within echo. This package was originally inspired from the
// https://github.com/ipfans/echo-session package, and modified to provide more
// functionality

package engine

import (
	"errors"
	"fmt"
	"log"

	"github.com/admpub/sessions"
	"github.com/webx-top/com"
	"github.com/webx-top/echo"
)

const (
	errorFormat = "[sessions] ERROR! %s\n"
)

var ErrInvalidSessionID = errors.New("invalid session ID")

type Session struct {
	name    string
	context echo.Context
	store   sessions.Store
	session *sessions.Session
	written bool
	preSave []func(echo.Context) error
}

func (s *Session) AddPreSaveHook(hook func(echo.Context) error) {
	s.preSave = append(s.preSave, hook)
}

func (s *Session) SetPreSaveHook(hooks ...func(echo.Context) error) {
	s.preSave = hooks
}

func (s *Session) Get(key string) interface{} {
	return s.Session().Values[key]
}

func (s *Session) Set(key string, val interface{}) echo.Sessioner {
	s.Session().Values[key] = val
	s.setWritten()
	return s
}

func (s *Session) Delete(key string) echo.Sessioner {
	delete(s.Session().Values, key)
	s.setWritten()
	return s
}

func (s *Session) Clear() echo.Sessioner {
	for key := range s.Session().Values {
		if k, ok := key.(string); ok {
			s.Delete(k)
		}
	}
	return s
}

func (s *Session) AddFlash(value interface{}, vars ...string) echo.Sessioner {
	s.Session().AddFlash(value, vars...)
	s.setWritten()
	return s
}

func (s *Session) Flashes(vars ...string) []interface{} {
	flashes := s.Session().Flashes(vars...)
	if len(flashes) > 0 {
		s.setWritten()
	}
	return flashes
}

func (s *Session) SetID(id string, notReload ...bool) error {
	if s.Session().ID == id {
		return nil
	}
	if !com.StrIsAlphaNumeric(id) {
		return ErrInvalidSessionID
	}
	s.Session().ID = id
	if len(notReload) == 0 || !notReload[0] {
		if err := s.Session().Reload(s.context); err != nil {
			return err
		}
		s.setWritten()
	}
	return nil
}

func (s *Session) ID() string {
	return s.Session().ID
}

func (s *Session) MustID() string {
	if len(s.Session().ID) > 0 {
		return s.Session().ID
	}
	if idGen, ok := s.Session().Store().(sessions.IDGenerator); ok {
		var err error
		s.Session().ID, err = idGen.GenerateID(s.context, s.Session())
		if err != nil {
			err = fmt.Errorf(`Session ID generation failed: %w`, err)
			panic(err)
		}
		return s.Session().ID
	}
	s.Session().ID = GenerateSessionID()
	return s.Session().ID
}

func (s *Session) RemoveID(sessionID string) error {
	if !com.StrIsAlphaNumeric(sessionID) {
		return ErrInvalidSessionID
	}
	return s.store.Remove(sessionID)
}

func (s *Session) Save() error {
	if !s.Written() {
		return nil
	}
	for _, hook := range s.preSave {
		if err := hook(s.context); err != nil {
			return err
		}
	}
	err := s.Session().Save(s.context)
	if err == nil {
		s.written = false
	} else {
		log.Printf(errorFormat, err)
	}
	return err
}

func (s *Session) Session() *sessions.Session {
	if s.session == nil {
		var err error
		s.session, err = s.store.Get(s.context, s.name)
		if err != nil {
			log.Printf(errorFormat, err)
		}
	}
	return s.session
}

func (s *Session) Written() bool {
	return s.written
}

func (s *Session) setWritten() *Session {
	s.written = true
	return s
}
<commit_msg>update<commit_after>// Session implements middleware for easily using github.com/gorilla/sessions
// within echo. This package was originally inspired from the
// https://github.com/ipfans/echo-session package, and modified to provide more
// functionality

package engine

import (
	"errors"
	"fmt"
	"log"

	"github.com/admpub/sessions"
	"github.com/webx-top/com"
	"github.com/webx-top/echo"
)

const (
	errorFormat = "[sessions] ERROR! %s\n"
)

var ErrInvalidSessionID = errors.New("invalid session ID")

type Session struct {
	name    string
	context echo.Context
	store   sessions.Store
	session *sessions.Session
	written bool
	preSave []func(echo.Context) error
}

func (s *Session) AddPreSaveHook(hook func(echo.Context) error) {
	s.preSave = append(s.preSave, hook)
}

func (s *Session) SetPreSaveHook(hooks ...func(echo.Context) error) {
	s.preSave = hooks
}

func (s *Session) Get(key string) interface{} {
	return s.Session().Values[key]
}

func (s *Session) Set(key string, val interface{}) echo.Sessioner {
	s.Session().Values[key] = val
	s.setWritten()
	return s
}

func (s *Session) Delete(key string) echo.Sessioner {
	delete(s.Session().Values, key)
	s.setWritten()
	return s
}

func (s *Session) Clear() echo.Sessioner {
	for key := range s.Session().Values {
		if k, ok := key.(string); ok {
			s.Delete(k)
		}
	}
	s.setWritten()
	return s
}

func (s *Session) AddFlash(value interface{}, vars ...string) echo.Sessioner {
	s.Session().AddFlash(value, vars...)
	s.setWritten()
	return s
}

func (s *Session) Flashes(vars ...string) []interface{} {
	flashes := s.Session().Flashes(vars...)
	if len(flashes) > 0 {
		s.setWritten()
	}
	return flashes
}

func (s *Session) SetID(id string, notReload ...bool) error {
	if s.Session().ID == id {
		return nil
	}
	if !com.StrIsAlphaNumeric(id) {
		return ErrInvalidSessionID
	}
	s.Session().ID = id
	if len(notReload) == 0 || !notReload[0] {
		if err := s.Session().Reload(s.context); err != nil {
			return err
		}
		s.setWritten()
	}
	return nil
}

func (s *Session) ID() string {
	return s.Session().ID
}

func (s *Session) MustID() string {
	if len(s.Session().ID) > 0 {
		return s.Session().ID
	}
	if idGen, ok := s.Session().Store().(sessions.IDGenerator); ok {
		var err error
		s.Session().ID, err = idGen.GenerateID(s.context, s.Session())
		if err != nil {
			err = fmt.Errorf(`Session ID generation failed: %w`, err)
			panic(err)
		}
		return s.Session().ID
	}
	s.Session().ID = GenerateSessionID()
	return s.Session().ID
}

func (s *Session) RemoveID(sessionID string) error {
	if !com.StrIsAlphaNumeric(sessionID) {
		return ErrInvalidSessionID
	}
	return s.store.Remove(sessionID)
}

func (s *Session) Save() error {
	if !s.Written() {
		return nil
	}
	for _, hook := range s.preSave {
		if err := hook(s.context); err != nil {
			return err
		}
	}
	err := s.Session().Save(s.context)
	if err == nil {
		s.written = false
	} else {
		log.Printf(errorFormat, err)
	}
	return err
}

func (s *Session) Session() *sessions.Session {
	if s.session == nil {
		var err error
		s.session, err = s.store.Get(s.context, s.name)
		if err != nil {
			log.Printf(errorFormat, err)
		}
	}
	return s.session
}

func (s *Session) Written() bool {
	return s.written
}

func (s *Session) setWritten() *Session {
	if !s.written {
		s.written = true
	}
	return s
}
<|endoftext|>


<commit_before>/* Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 clownfish

import "testing"
import "reflect"

func TestByteBufCat(t *testing.T) {
	bb := NewByteBuf(0)
	content := []byte("foo")
	bb.cat(content)
	if got := bb.yieldBlob(); !reflect.DeepEqual(got, content) {
		t.Errorf("Expected %v, got %v", content, got)
	}
}

func TestByteBufSetSizeGetSize(t *testing.T) {
	bb := NewByteBuf(0)
	content := []byte("abc")
	bb.cat(content)
	bb.setSize(2)
	if got := bb.getSize(); got != 2 {
		t.Errorf("Expected size 2, got %d", got)
	}
	expected := []byte("ab")
	if got := bb.yieldBlob(); !reflect.DeepEqual(got, expected) {
		t.Errorf("Expected %v, got %v", expected, got)
	}
}

func TestByteBufGetCapacity(t *testing.T) {
	bb := NewByteBuf(5)
	if cap := bb.getCapacity(); cap < 5 {
		t.Errorf("Expected at least 5, got %d", cap)
	}
}

func TestByteBufEquals(t *testing.T) {
	bb := NewByteBuf(0)
	other := NewByteBuf(0)
	content := []byte("foo")
	bb.cat(content)
	other.cat(content)
	if !bb.Equals(other) {
		t.Errorf("Equals against equal ByteBuf")
	}
	other.setSize(2)
	if bb.Equals(other) {
		t.Errorf("Equals against non-equal ByteBuf")
	}
	if bb.Equals(42) {
		t.Errorf("Equals against arbitrary Go type")
	}
}

func TestByteBufClone(t *testing.T) {
	content := []byte("foo")
	bb := NewByteBuf(0)
	bb.cat(content)
	clone := bb.Clone().(ByteBuf)
	if got := clone.yieldBlob(); !reflect.DeepEqual(got, content) {
		t.Errorf("Expected %v, got %v", content, got)
	}
}

func TestByteBufCompareTo(t *testing.T) {
	bb := NewByteBuf(0)
	other := NewByteBuf(0)
	content := []byte("foo")
	bb.cat(content)
	other.cat(content)
	if got := bb.CompareTo(other); got != 0 {
		t.Errorf("CompareTo equal, got %d", got)
	}
	other.setSize(2)
	if got := bb.CompareTo(other); got <= 0 {
		t.Errorf("CompareTo lesser, got %d", got)
	}
	if got := other.CompareTo(bb); got >= 0 {
		t.Errorf("CompareTo greater, got %d", got)
	}
}
<commit_msg>Fix public ByteBuf methods in Go bindings<commit_after>/* Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 clownfish

import "testing"
import "reflect"

func TestByteBufCat(t *testing.T) {
	bb := NewByteBuf(0)
	content := []byte("foo")
	bb.Cat(content)
	if got := bb.YieldBlob(); !reflect.DeepEqual(got, content) {
		t.Errorf("Expected %v, got %v", content, got)
	}
}

func TestByteBufSetSizeGetSize(t *testing.T) {
	bb := NewByteBuf(0)
	content := []byte("abc")
	bb.Cat(content)
	bb.SetSize(2)
	if got := bb.GetSize(); got != 2 {
		t.Errorf("Expected size 2, got %d", got)
	}
	expected := []byte("ab")
	if got := bb.YieldBlob(); !reflect.DeepEqual(got, expected) {
		t.Errorf("Expected %v, got %v", expected, got)
	}
}

func TestByteBufGetCapacity(t *testing.T) {
	bb := NewByteBuf(5)
	if cap := bb.GetCapacity(); cap < 5 {
		t.Errorf("Expected at least 5, got %d", cap)
	}
}

func TestByteBufEquals(t *testing.T) {
	bb := NewByteBuf(0)
	other := NewByteBuf(0)
	content := []byte("foo")
	bb.Cat(content)
	other.Cat(content)
	if !bb.Equals(other) {
		t.Errorf("Equals against equal ByteBuf")
	}
	other.SetSize(2)
	if bb.Equals(other) {
		t.Errorf("Equals against non-equal ByteBuf")
	}
	if bb.Equals(42) {
		t.Errorf("Equals against arbitrary Go type")
	}
}

func TestByteBufClone(t *testing.T) {
	content := []byte("foo")
	bb := NewByteBuf(0)
	bb.Cat(content)
	clone := bb.Clone().(ByteBuf)
	if got := clone.YieldBlob(); !reflect.DeepEqual(got, content) {
		t.Errorf("Expected %v, got %v", content, got)
	}
}

func TestByteBufCompareTo(t *testing.T) {
	bb := NewByteBuf(0)
	other := NewByteBuf(0)
	content := []byte("foo")
	bb.Cat(content)
	other.Cat(content)
	if got := bb.CompareTo(other); got != 0 {
		t.Errorf("CompareTo equal, got %d", got)
	}
	other.SetSize(2)
	if got := bb.CompareTo(other); got <= 0 {
		t.Errorf("CompareTo lesser, got %d", got)
	}
	if got := other.CompareTo(bb); got >= 0 {
		t.Errorf("CompareTo greater, got %d", got)
	}
}
<|endoftext|>


<commit_before>package collector

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"path"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
)

type snapshotMetric struct {
	Type   prometheus.ValueType
	Desc   *prometheus.Desc
	Value  func(snapshotStats SnapshotStatDataResponse) float64
	Labels func(repositoryName string, snapshotStats SnapshotStatDataResponse) []string
}

type repositoryMetric struct {
	Type   prometheus.ValueType
	Desc   *prometheus.Desc
	Value  func(snapshotsStats SnapshotStatsResponse) float64
	Labels func(repositoryName string) []string
}

var (
	defaultSnapshotLabels      = []string{"repository", "state", "version"}
	defaultSnapshotLabelValues = func(repositoryName string, snapshotStats SnapshotStatDataResponse) []string {
		return []string{repositoryName, snapshotStats.State, snapshotStats.Version}
	}
	defaultSnapshotRepositoryLabels      = []string{"repository"}
	defaultSnapshotRepositoryLabelValues = func(repositoryName string) []string {
		return []string{repositoryName}
	}
)

// Snapshots information struct
type Snapshots struct {
	logger log.Logger
	client *http.Client
	url    *url.URL

	up                              prometheus.Gauge
	totalScrapes, jsonParseFailures prometheus.Counter

	snapshotMetrics   []*snapshotMetric
	repositoryMetrics []*repositoryMetric
}

// NewSnapshots defines Snapshots Prometheus metrics
func NewSnapshots(logger log.Logger, client *http.Client, url *url.URL) *Snapshots {
	return &Snapshots{
		logger: logger,
		client: client,
		url:    url,

		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Name: prometheus.BuildFQName(namespace, "snapshot_stats", "up"),
			Help: "Was the last scrape of the ElasticSearch snapshots endpoint successful.",
		}),
		totalScrapes: prometheus.NewCounter(prometheus.CounterOpts{
			Name: prometheus.BuildFQName(namespace, "snapshot_stats", "total_scrapes"),
			Help: "Current total ElasticSearch snapshots scrapes.",
		}),
		jsonParseFailures: prometheus.NewCounter(prometheus.CounterOpts{
			Name: prometheus.BuildFQName(namespace, "snapshot_stats", "json_parse_failures"),
			Help: "Number of errors while parsing JSON.",
		}),
		snapshotMetrics: []*snapshotMetric{
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_number_of_indices"),
					"Number of indices in the last snapshot",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(len(snapshotStats.Indices))
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_start_time_timestamp"),
					"Last snapshot start timestamp",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.StartTimeInMillis / 1000)
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_end_time_timestamp"),
					"Last snapshot end timestamp",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.EndTimeInMillis / 1000)
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_number_of_failures"),
					"Last snapshot number of failures",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(len(snapshotStats.Failures))
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_total_shards"),
					"Last snapshot total shards",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.Shards.Total)
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_failed_shards"),
					"Last snapshot failed shards",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.Shards.Failed)
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_successful_shards"),
					"Last snapshot successful shards",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.Shards.Successful)
				},
				Labels: defaultSnapshotLabelValues,
			},
		},
		repositoryMetrics: []*repositoryMetric{
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "number_of_snapshots"),
					"Number of snapshots in a repository",
					defaultSnapshotRepositoryLabels, nil,
				),
				Value: func(snapshotsStats SnapshotStatsResponse) float64 {
					return float64(len(snapshotsStats.Snapshots))
				},
				Labels: defaultSnapshotRepositoryLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "oldest_snapshot_timestamp"),
					"Timestamp of the oldest snapshot",
					defaultSnapshotRepositoryLabels, nil,
				),
				Value: func(snapshotsStats SnapshotStatsResponse) float64 {
					if len(snapshotsStats.Snapshots) == 0 {
						return 0
					}
					return float64(snapshotsStats.Snapshots[0].StartTimeInMillis / 1000)
				},
				Labels: defaultSnapshotRepositoryLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "latest_snapshot_timestamp"),
					"Timestamp of the latest SUCCESS or PARTIAL snapshot",
					defaultSnapshotRepositoryLabels, nil,
				),
				Value: func(snapshotsStats SnapshotStatsResponse) float64 {
					for i := len(snapshotsStats.Snapshots) - 1; i >= 0; i-- {
						var snap = snapshotsStats.Snapshots[i]
						if snap.State == "SUCCESS" || snap.State == "PARTIAL" {
							return float64(snap.StartTimeInMillis / 1000)
						}
					}
					return 0
				},
				Labels: defaultSnapshotRepositoryLabelValues,
			},
		},
	}
}

// Describe add Snapshots metrics descriptions
func (s *Snapshots) Describe(ch chan<- *prometheus.Desc) {
	for _, metric := range s.snapshotMetrics {
		ch <- metric.Desc
	}
	ch <- s.up.Desc()
	ch <- s.totalScrapes.Desc()
	ch <- s.jsonParseFailures.Desc()
}

func (s *Snapshots) getAndParseURL(u *url.URL, data interface{}) error {
	res, err := s.client.Get(u.String())
	if err != nil {
		return fmt.Errorf("failed to get from %s://%s:%s%s: %s",
			u.Scheme, u.Hostname(), u.Port(), u.Path, err)
	}

	defer func() {
		err = res.Body.Close()
		if err != nil {
			_ = level.Warn(s.logger).Log(
				"msg", "failed to close http.Client",
				"err", err,
			)
		}
	}()

	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP Request failed with code %d", res.StatusCode)
	}

	if err := json.NewDecoder(res.Body).Decode(data); err != nil {
		s.jsonParseFailures.Inc()
		return err
	}
	return nil
}

func (s *Snapshots) fetchAndDecodeSnapshotsStats() (map[string]SnapshotStatsResponse, error) {
	mssr := make(map[string]SnapshotStatsResponse)

	u := *s.url
	u.Path = path.Join(u.Path, "/_snapshot")
	var srr SnapshotRepositoriesResponse
	err := s.getAndParseURL(&u, &srr)
	if err != nil {
		return nil, err
	}
	for repository := range srr {
		u := *s.url
		u.Path = path.Join(u.Path, "/_snapshot", repository, "/_all")
		var ssr SnapshotStatsResponse
		err := s.getAndParseURL(&u, &ssr)
		if err != nil {
			continue
		}
		mssr[repository] = ssr
	}

	return mssr, nil
}

// Collect gets Snapshots metric values
func (s *Snapshots) Collect(ch chan<- prometheus.Metric) {
	s.totalScrapes.Inc()
	defer func() {
		ch <- s.up
		ch <- s.totalScrapes
		ch <- s.jsonParseFailures
	}()

	// indices
	snapshotsStatsResp, err := s.fetchAndDecodeSnapshotsStats()
	if err != nil {
		s.up.Set(0)
		_ = level.Warn(s.logger).Log(
			"msg", "failed to fetch and decode snapshot stats",
			"err", err,
		)
		return
	}
	s.up.Set(1)

	// Snapshots stats
	for repositoryName, snapshotStats := range snapshotsStatsResp {
		for _, metric := range s.repositoryMetrics {
			ch <- prometheus.MustNewConstMetric(
				metric.Desc,
				metric.Type,
				metric.Value(snapshotStats),
				metric.Labels(repositoryName)...,
			)
		}
		if len(snapshotStats.Snapshots) == 0 {
			continue
		}

		lastSnapshot := snapshotStats.Snapshots[len(snapshotStats.Snapshots)-1]
		for _, metric := range s.snapshotMetrics {
			ch <- prometheus.MustNewConstMetric(
				metric.Desc,
				metric.Type,
				metric.Value(lastSnapshot),
				metric.Labels(repositoryName, lastSnapshot)...,
			)
		}
	}
}
<commit_msg>fix: metric name<commit_after>package collector

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"path"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
)

type snapshotMetric struct {
	Type   prometheus.ValueType
	Desc   *prometheus.Desc
	Value  func(snapshotStats SnapshotStatDataResponse) float64
	Labels func(repositoryName string, snapshotStats SnapshotStatDataResponse) []string
}

type repositoryMetric struct {
	Type   prometheus.ValueType
	Desc   *prometheus.Desc
	Value  func(snapshotsStats SnapshotStatsResponse) float64
	Labels func(repositoryName string) []string
}

var (
	defaultSnapshotLabels      = []string{"repository", "state", "version"}
	defaultSnapshotLabelValues = func(repositoryName string, snapshotStats SnapshotStatDataResponse) []string {
		return []string{repositoryName, snapshotStats.State, snapshotStats.Version}
	}
	defaultSnapshotRepositoryLabels      = []string{"repository"}
	defaultSnapshotRepositoryLabelValues = func(repositoryName string) []string {
		return []string{repositoryName}
	}
)

// Snapshots information struct
type Snapshots struct {
	logger log.Logger
	client *http.Client
	url    *url.URL

	up                              prometheus.Gauge
	totalScrapes, jsonParseFailures prometheus.Counter

	snapshotMetrics   []*snapshotMetric
	repositoryMetrics []*repositoryMetric
}

// NewSnapshots defines Snapshots Prometheus metrics
func NewSnapshots(logger log.Logger, client *http.Client, url *url.URL) *Snapshots {
	return &Snapshots{
		logger: logger,
		client: client,
		url:    url,

		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Name: prometheus.BuildFQName(namespace, "snapshot_stats", "up"),
			Help: "Was the last scrape of the ElasticSearch snapshots endpoint successful.",
		}),
		totalScrapes: prometheus.NewCounter(prometheus.CounterOpts{
			Name: prometheus.BuildFQName(namespace, "snapshot_stats", "total_scrapes"),
			Help: "Current total ElasticSearch snapshots scrapes.",
		}),
		jsonParseFailures: prometheus.NewCounter(prometheus.CounterOpts{
			Name: prometheus.BuildFQName(namespace, "snapshot_stats", "json_parse_failures"),
			Help: "Number of errors while parsing JSON.",
		}),
		snapshotMetrics: []*snapshotMetric{
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_number_of_indices"),
					"Number of indices in the last snapshot",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(len(snapshotStats.Indices))
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_start_time_timestamp"),
					"Last snapshot start timestamp",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.StartTimeInMillis / 1000)
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_end_time_timestamp"),
					"Last snapshot end timestamp",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.EndTimeInMillis / 1000)
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_number_of_failures"),
					"Last snapshot number of failures",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(len(snapshotStats.Failures))
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_total_shards"),
					"Last snapshot total shards",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.Shards.Total)
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_failed_shards"),
					"Last snapshot failed shards",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.Shards.Failed)
				},
				Labels: defaultSnapshotLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "snapshot_successful_shards"),
					"Last snapshot successful shards",
					defaultSnapshotLabels, nil,
				),
				Value: func(snapshotStats SnapshotStatDataResponse) float64 {
					return float64(snapshotStats.Shards.Successful)
				},
				Labels: defaultSnapshotLabelValues,
			},
		},
		repositoryMetrics: []*repositoryMetric{
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "number_of_snapshots"),
					"Number of snapshots in a repository",
					defaultSnapshotRepositoryLabels, nil,
				),
				Value: func(snapshotsStats SnapshotStatsResponse) float64 {
					return float64(len(snapshotsStats.Snapshots))
				},
				Labels: defaultSnapshotRepositoryLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "oldest_snapshot_timestamp"),
					"Timestamp of the oldest snapshot",
					defaultSnapshotRepositoryLabels, nil,
				),
				Value: func(snapshotsStats SnapshotStatsResponse) float64 {
					if len(snapshotsStats.Snapshots) == 0 {
						return 0
					}
					return float64(snapshotsStats.Snapshots[0].StartTimeInMillis / 1000)
				},
				Labels: defaultSnapshotRepositoryLabelValues,
			},
			{
				Type: prometheus.GaugeValue,
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, "snapshot_stats", "latest_snapshot_seconds"),
					"Timestamp of the latest SUCCESS or PARTIAL snapshot",
					defaultSnapshotRepositoryLabels, nil,
				),
				Value: func(snapshotsStats SnapshotStatsResponse) float64 {
					for i := len(snapshotsStats.Snapshots) - 1; i >= 0; i-- {
						var snap = snapshotsStats.Snapshots[i]
						if snap.State == "SUCCESS" || snap.State == "PARTIAL" {
							return float64(snap.StartTimeInMillis / 1000)
						}
					}
					return 0
				},
				Labels: defaultSnapshotRepositoryLabelValues,
			},
		},
	}
}

// Describe add Snapshots metrics descriptions
func (s *Snapshots) Describe(ch chan<- *prometheus.Desc) {
	for _, metric := range s.snapshotMetrics {
		ch <- metric.Desc
	}
	ch <- s.up.Desc()
	ch <- s.totalScrapes.Desc()
	ch <- s.jsonParseFailures.Desc()
}

func (s *Snapshots) getAndParseURL(u *url.URL, data interface{}) error {
	res, err := s.client.Get(u.String())
	if err != nil {
		return fmt.Errorf("failed to get from %s://%s:%s%s: %s",
			u.Scheme, u.Hostname(), u.Port(), u.Path, err)
	}

	defer func() {
		err = res.Body.Close()
		if err != nil {
			_ = level.Warn(s.logger).Log(
				"msg", "failed to close http.Client",
				"err", err,
			)
		}
	}()

	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP Request failed with code %d", res.StatusCode)
	}

	if err := json.NewDecoder(res.Body).Decode(data); err != nil {
		s.jsonParseFailures.Inc()
		return err
	}
	return nil
}

func (s *Snapshots) fetchAndDecodeSnapshotsStats() (map[string]SnapshotStatsResponse, error) {
	mssr := make(map[string]SnapshotStatsResponse)

	u := *s.url
	u.Path = path.Join(u.Path, "/_snapshot")
	var srr SnapshotRepositoriesResponse
	err := s.getAndParseURL(&u, &srr)
	if err != nil {
		return nil, err
	}
	for repository := range srr {
		u := *s.url
		u.Path = path.Join(u.Path, "/_snapshot", repository, "/_all")
		var ssr SnapshotStatsResponse
		err := s.getAndParseURL(&u, &ssr)
		if err != nil {
			continue
		}
		mssr[repository] = ssr
	}

	return mssr, nil
}

// Collect gets Snapshots metric values
func (s *Snapshots) Collect(ch chan<- prometheus.Metric) {
	s.totalScrapes.Inc()
	defer func() {
		ch <- s.up
		ch <- s.totalScrapes
		ch <- s.jsonParseFailures
	}()

	// indices
	snapshotsStatsResp, err := s.fetchAndDecodeSnapshotsStats()
	if err != nil {
		s.up.Set(0)
		_ = level.Warn(s.logger).Log(
			"msg", "failed to fetch and decode snapshot stats",
			"err", err,
		)
		return
	}
	s.up.Set(1)

	// Snapshots stats
	for repositoryName, snapshotStats := range snapshotsStatsResp {
		for _, metric := range s.repositoryMetrics {
			ch <- prometheus.MustNewConstMetric(
				metric.Desc,
				metric.Type,
				metric.Value(snapshotStats),
				metric.Labels(repositoryName)...,
			)
		}
		if len(snapshotStats.Snapshots) == 0 {
			continue
		}

		lastSnapshot := snapshotStats.Snapshots[len(snapshotStats.Snapshots)-1]
		for _, metric := range s.snapshotMetrics {
			ch <- prometheus.MustNewConstMetric(
				metric.Desc,
				metric.Type,
				metric.Value(lastSnapshot),
				metric.Labels(repositoryName, lastSnapshot)...,
			)
		}
	}
}
<|endoftext|>