<commit_before>package main

// See https://github.com/grafana/grafana/blob/master/docs/sources/developers/plugins/backend.md for
// details on grafana backend plugins

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"net/url"
	"reflect"
	"runtime/debug"
	"strconv"
	"strings"

	"github.com/bitly/go-simplejson"
	"github.com/grafana/grafana-plugin-model/go/datasource"
	"github.com/hashicorp/go-plugin"
	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

var httpClient = &http.Client{}

type ClickhouseDatasource struct {
	plugin.NetRPCUnsupportedPlugin
}

func (t *ClickhouseDatasource) Query(ctx context.Context, req *datasource.DatasourceRequest) (r *datasource.DatasourceResponse, err error) {
	// catch all panics and override err return value
	defer func() {
		if panicMsg := recover(); panicMsg != nil {
			err = fmt.Errorf("clickhouse plugin panicked: %+v stacktrace:\n%s", panicMsg, debug.Stack())
		}
	}()

	refId := req.Queries[0].RefId
	modelJson, err := simplejson.NewJson([]byte(req.Queries[0].ModelJson))
	if err != nil {
		return nil, fmt.Errorf("unable to parse query: %w", err)
	}

	query := modelJson.Get("rawQuery").MustString()
	request, err := createRequest(req, query)
	if err != nil {
		return nil, err
	}

	response, err := ctxhttp.Do(ctx, httpClient, request)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := response.Body.Close(); err!=nil {
			log.Fatal("can't close HTTP Response body")
		}
	}()

	// Body must be drained and closed on each request as per the docs: https://golang.org/pkg/net/http/#Client.Do
	// otherwise the http client connection cannot be reused
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	if response.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("invalid status code. status: %v", response.Status)
	}

	return parseResponse(body, refId)
}

func createRequest(req *datasource.DatasourceRequest, query string) (*http.Request, error) {
	body := ""
	method := http.MethodGet
	headers := http.Header{}
	dataSourceUrl, err := url.Parse(req.Datasource.Url)
	if err != nil {
		return nil, fmt.Errorf("unable to parse clickhouse dataSourceUrl: %w", err)
	}

	params := dataSourceUrl.Query()
	params.Add("query", query+" FORMAT JSON")

	/*
	 Note: The current plugins model does not support basic authorization.
	 We have access to basicAuthPassword but not the basic auth name. Users
	 will have to use the useYandexCloudAuthorization
	 option instead for clickhouse auth.
	 This will be necessary until the new grafana plugin model becomes available:
	 https://github.com/grafana/grafana-plugin-sdk-go
	*/
	secureOptions := req.Datasource.DecryptedSecureJsonData
	options := make(map[string]interface{})
	err = json.Unmarshal([]byte(req.Datasource.JsonData), &options)
	if err != nil {
		return nil, fmt.Errorf("unable to parse clickhouse options: %w", err)
	}

	for k, v := range options {
		switch k {
		case "usePOST":
			method = http.MethodPost
			params.Del("query")
			body = query
			break
		case "defaultDatabase":
			db, _ := v.(string)
			params.Add("database", db)
			break
		case "addCorsHeaders":
			params.Add("add_http_cors_header", "1")
			break
		case "useYandexCloudAuthorization":
			if user, ok := options["xHeaderUser"]; ok {
				chUser, _ := user.(string)
				headers.Add("X-ClickHouse-User", chUser)
			}

			if key, ok := options["xHeaderKey"]; ok {
				chKey, _ := key.(string)
				headers.Add("X-ClickHouse-Key", chKey)
			}
			break
		default:
			if strings.HasPrefix(k, "httpHeaderName") {
				headerKey := strings.Replace(k, "Name", "Value", 1)
				value := ""
				name, _ := v.(string)
				if hv, ok := secureOptions[headerKey]; ok {
					value = hv
				}

				headers.Add(name, value)
			}
			break
		}
	}

	dataSourceUrl.RawQuery = params.Encode()
	request, err := http.NewRequest(method, dataSourceUrl.String(), bytes.NewBufferString(body))
	if err != nil {
		return nil, err
	}

	request.Header = headers
	return request, nil
}

var floatType = reflect.TypeOf(float64(0))
var stringType = reflect.TypeOf("")

func parseFloat64(v interface{}) (float64, error) {
	switch i := v.(type) {
	case string:
		return strconv.ParseFloat(i, 64)
	case float64:
		return i, nil
	case float32:
		return float64(i), nil
	case int64:
		return float64(i), nil
	case int32:
		return float64(i), nil
	case int:
		return float64(i), nil
	case uint64:
		return float64(i), nil
	case uint32:
		return float64(i), nil
	case uint:
		return float64(i), nil
	default:
		tv := reflect.ValueOf(i)
		tv = reflect.Indirect(tv)
		if tv.Type().ConvertibleTo(floatType) {
			fv := tv.Convert(floatType)
			return fv.Float(), nil
		} else if tv.Type().ConvertibleTo(stringType) {
			sv := tv.Convert(stringType)
			s := sv.String()
			return strconv.ParseFloat(s, 64)
		} else {
			return math.NaN(), fmt.Errorf("can't convert %v to float64", tv.Type())
		}
	}
}

func parseResponse(body []byte, refId string) (*datasource.DatasourceResponse, error) {

	parsedBody := ClickHouseResponse{}
	err := json.Unmarshal(body, &parsedBody)
	if err != nil {
		return nil, fmt.Errorf("unable to parse response body: %s\n\n parsing error: %w", body, err)
	}

	seriesMap := map[string]*datasource.TimeSeries{}
	metaTypesMap := map[string]string{}
	// expect first column as timestamp
	tsMetaName := parsedBody.Meta[0].Name
	for _, meta := range parsedBody.Meta {
		if meta.Name != tsMetaName && !strings.HasPrefix(meta.Type,"Array(Tuple("){
			seriesMap[meta.Name] = &datasource.TimeSeries{Name: meta.Name, Points: []*datasource.Point{}}
		}
		metaTypesMap[meta.Name] = meta.Type
	}

	for _, dataPoint := range parsedBody.Data {
		timestamp, err := strconv.ParseInt(dataPoint[tsMetaName].(string), 10, 64)
		if err != nil {
			return nil, fmt.Errorf("unable to parse timestamp with alias=`%s` value=%s error=%w", tsMetaName, dataPoint[tsMetaName].(string), err)
		}
		for k, v := range dataPoint {
			if k != tsMetaName {
				var point float64
				var err error

				if !strings.HasPrefix(metaTypesMap[k],"Array(Tuple(") {

					point, err = parseFloat64(v)
					if err != nil {
						return nil, fmt.Errorf("unable to parse value %v for '%s': %w", v, k, err)
					}

					seriesMap[k].Points = append(seriesMap[k].Points, &datasource.Point{
						Timestamp: timestamp,
						Value:     point,
					})
				} else {
					var arrayOfTuples [][]string
					switch arrays := v.(type) {
					case []interface{}:
						for _, array := range arrays {
							switch tuple := array.(type) {
							case []interface{}:
								var t []string
								for _, s := range tuple {
									t = append(t, fmt.Sprintf("%v",s))
								}
								arrayOfTuples = append(arrayOfTuples, t)
							default:
								return nil, fmt.Errorf("unable to parse data section type=%T in response json: %s", tuple, tuple)
							}
						}
					default:
						return nil, fmt.Errorf("unable to parse data section type=%T in response json: %s", v, v)
					}
					for _, tuple := range arrayOfTuples {
						tsName := tuple[0]
						tsValue := tuple[1]
						ts, isExists := seriesMap[tsName]
						if !isExists {
							ts = &datasource.TimeSeries{Name: tsName, Points: []*datasource.Point{}}
							seriesMap[tsName] = ts
						}
						point, err = parseFloat64(tsValue)
						if err != nil {
							return nil, fmt.Errorf("unable to parse value %v for '%s': %w", tsValue, tsName, err)
						}
						ts.Points = append(ts.Points, &datasource.Point{
							Timestamp: timestamp,
							Value: point,
						})
					}
				}
			}
		}
	}

	var series []*datasource.TimeSeries
	for _, timeSeries := range seriesMap {
		series = append(series, timeSeries)
	}

	metaJSON, _ := json.Marshal(parsedBody.Meta)
	return &datasource.DatasourceResponse{
		Results: []*datasource.QueryResult{
			{
				Series:   series,
				RefId:    refId,
				MetaJson: string(metaJSON),
			},
		},
	}, nil
}

type ClickHouseResponse struct {
	Meta []ClickHouseMeta
	Data []map[string]interface{}
	Rows int
}

type ClickHouseMeta struct {
	Name string
	Type string
}
<commit_msg>make query from backend in JSON format instead TSV (alerting support need JSON)<commit_after>package main

// See https://github.com/grafana/grafana/blob/master/docs/sources/developers/plugins/backend.md for
// details on grafana backend plugins

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"net/url"
	"reflect"
	"runtime/debug"
	"strconv"
	"strings"

	"github.com/bitly/go-simplejson"
	"github.com/grafana/grafana-plugin-model/go/datasource"
	"github.com/hashicorp/go-plugin"
	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

var httpClient = &http.Client{}

type ClickhouseDatasource struct {
	plugin.NetRPCUnsupportedPlugin
}

func (t *ClickhouseDatasource) Query(ctx context.Context, req *datasource.DatasourceRequest) (r *datasource.DatasourceResponse, err error) {
	// catch all panics and override err return value
	defer func() {
		if panicMsg := recover(); panicMsg != nil {
			err = fmt.Errorf("clickhouse plugin panicked: %+v stacktrace:\n%s", panicMsg, debug.Stack())
		}
	}()

	refId := req.Queries[0].RefId
	modelJson, err := simplejson.NewJson([]byte(req.Queries[0].ModelJson))
	if err != nil {
		return nil, fmt.Errorf("unable to parse query: %w", err)
	}

	query := modelJson.Get("rawQuery").MustString()
	request, err := createRequest(req, query)
	if err != nil {
		return nil, err
	}

	response, err := ctxhttp.Do(ctx, httpClient, request)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := response.Body.Close(); err!=nil {
			log.Fatal("can't close HTTP Response body")
		}
	}()

	// Body must be drained and closed on each request as per the docs: https://golang.org/pkg/net/http/#Client.Do
	// otherwise the http client connection cannot be reused
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	if response.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("invalid status code. status: %v", response.Status)
	}

	return parseResponse(body, refId)
}

func createRequest(req *datasource.DatasourceRequest, query string) (*http.Request, error) {
	body := ""
	method := http.MethodGet
	headers := http.Header{}
	dataSourceUrl, err := url.Parse(req.Datasource.Url)
	if err != nil {
		return nil, fmt.Errorf("unable to parse clickhouse dataSourceUrl: %w", err)
	}

	params := dataSourceUrl.Query()
	params.Add("query", query+" FORMAT JSON")

	/*
	 Note: The current plugins model does not support basic authorization.
	 We have access to basicAuthPassword but not the basic auth name. Users
	 will have to use the useYandexCloudAuthorization
	 option instead for clickhouse auth.
	 This will be necessary until the new grafana plugin model becomes available:
	 https://github.com/grafana/grafana-plugin-sdk-go
	*/
	secureOptions := req.Datasource.DecryptedSecureJsonData
	options := make(map[string]interface{})
	err = json.Unmarshal([]byte(req.Datasource.JsonData), &options)
	if err != nil {
		return nil, fmt.Errorf("unable to parse clickhouse options: %w", err)
	}

	for k, v := range options {
		switch k {
		case "usePOST":
			method = http.MethodPost
			params.Del("query")
			body = query+" FORMAT JSON"
			break
		case "defaultDatabase":
			db, _ := v.(string)
			params.Add("database", db)
			break
		case "addCorsHeaders":
			params.Add("add_http_cors_header", "1")
			break
		case "useYandexCloudAuthorization":
			if user, ok := options["xHeaderUser"]; ok {
				chUser, _ := user.(string)
				headers.Add("X-ClickHouse-User", chUser)
			}

			if key, ok := options["xHeaderKey"]; ok {
				chKey, _ := key.(string)
				headers.Add("X-ClickHouse-Key", chKey)
			}
			break
		default:
			if strings.HasPrefix(k, "httpHeaderName") {
				headerKey := strings.Replace(k, "Name", "Value", 1)
				value := ""
				name, _ := v.(string)
				if hv, ok := secureOptions[headerKey]; ok {
					value = hv
				}

				headers.Add(name, value)
			}
			break
		}
	}

	dataSourceUrl.RawQuery = params.Encode()
	request, err := http.NewRequest(method, dataSourceUrl.String(), bytes.NewBufferString(body))
	if err != nil {
		return nil, err
	}

	request.Header = headers
	return request, nil
}

var floatType = reflect.TypeOf(float64(0))
var stringType = reflect.TypeOf("")

func parseFloat64(v interface{}) (float64, error) {
	switch i := v.(type) {
	case string:
		return strconv.ParseFloat(i, 64)
	case float64:
		return i, nil
	case float32:
		return float64(i), nil
	case int64:
		return float64(i), nil
	case int32:
		return float64(i), nil
	case int:
		return float64(i), nil
	case uint64:
		return float64(i), nil
	case uint32:
		return float64(i), nil
	case uint:
		return float64(i), nil
	default:
		tv := reflect.ValueOf(i)
		tv = reflect.Indirect(tv)
		if tv.Type().ConvertibleTo(floatType) {
			fv := tv.Convert(floatType)
			return fv.Float(), nil
		} else if tv.Type().ConvertibleTo(stringType) {
			sv := tv.Convert(stringType)
			s := sv.String()
			return strconv.ParseFloat(s, 64)
		} else {
			return math.NaN(), fmt.Errorf("can't convert %v to float64", tv.Type())
		}
	}
}

func parseResponse(body []byte, refId string) (*datasource.DatasourceResponse, error) {

	parsedBody := ClickHouseResponse{}
	err := json.Unmarshal(body, &parsedBody)
	if err != nil {
		return nil, fmt.Errorf("unable to parse response body: %s\n\n parsing error: %w", body, err)
	}

	seriesMap := map[string]*datasource.TimeSeries{}
	metaTypesMap := map[string]string{}
	// expect first column as timestamp
	tsMetaName := parsedBody.Meta[0].Name
	for _, meta := range parsedBody.Meta {
		if meta.Name != tsMetaName && !strings.HasPrefix(meta.Type,"Array(Tuple("){
			seriesMap[meta.Name] = &datasource.TimeSeries{Name: meta.Name, Points: []*datasource.Point{}}
		}
		metaTypesMap[meta.Name] = meta.Type
	}

	for _, dataPoint := range parsedBody.Data {
		timestamp, err := strconv.ParseInt(dataPoint[tsMetaName].(string), 10, 64)
		if err != nil {
			return nil, fmt.Errorf("unable to parse timestamp with alias=`%s` value=%s error=%w", tsMetaName, dataPoint[tsMetaName].(string), err)
		}
		for k, v := range dataPoint {
			if k != tsMetaName {
				var point float64
				var err error

				if !strings.HasPrefix(metaTypesMap[k],"Array(Tuple(") {

					point, err = parseFloat64(v)
					if err != nil {
						return nil, fmt.Errorf("unable to parse value %v for '%s': %w", v, k, err)
					}

					seriesMap[k].Points = append(seriesMap[k].Points, &datasource.Point{
						Timestamp: timestamp,
						Value:     point,
					})
				} else {
					var arrayOfTuples [][]string
					switch arrays := v.(type) {
					case []interface{}:
						for _, array := range arrays {
							switch tuple := array.(type) {
							case []interface{}:
								var t []string
								for _, s := range tuple {
									t = append(t, fmt.Sprintf("%v",s))
								}
								arrayOfTuples = append(arrayOfTuples, t)
							default:
								return nil, fmt.Errorf("unable to parse data section type=%T in response json: %s", tuple, tuple)
							}
						}
					default:
						return nil, fmt.Errorf("unable to parse data section type=%T in response json: %s", v, v)
					}
					for _, tuple := range arrayOfTuples {
						tsName := tuple[0]
						tsValue := tuple[1]
						ts, isExists := seriesMap[tsName]
						if !isExists {
							ts = &datasource.TimeSeries{Name: tsName, Points: []*datasource.Point{}}
							seriesMap[tsName] = ts
						}
						point, err = parseFloat64(tsValue)
						if err != nil {
							return nil, fmt.Errorf("unable to parse value %v for '%s': %w", tsValue, tsName, err)
						}
						ts.Points = append(ts.Points, &datasource.Point{
							Timestamp: timestamp,
							Value: point,
						})
					}
				}
			}
		}
	}

	var series []*datasource.TimeSeries
	for _, timeSeries := range seriesMap {
		series = append(series, timeSeries)
	}

	metaJSON, _ := json.Marshal(parsedBody.Meta)
	return &datasource.DatasourceResponse{
		Results: []*datasource.QueryResult{
			{
				Series:   series,
				RefId:    refId,
				MetaJson: string(metaJSON),
			},
		},
	}, nil
}

type ClickHouseResponse struct {
	Meta []ClickHouseMeta
	Data []map[string]interface{}
	Rows int
}

type ClickHouseMeta struct {
	Name string
	Type string
}
<|endoftext|>


<commit_before>// Copyright 2014 Jarmo Puttonen <jarmo.puttonen@gmail.com>. All rights reserved.
// Use of this source code is governed by a MIT-style
// licence that can be found in the LICENCE file.

/*Package paparazzogo implements a caching proxy for
serving MJPEG-stream as JPG-images.
*/
package paparazzogo

import (
	"bytes"
	"errors"
	"io"
	"log"
	"mime"
	"mime/multipart"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// A Mjpegproxy implements http.Handler	interface and generates
// JPG-images from a MJPEG-stream.
type Mjpegproxy struct {
	partbufsize      int64
	waittime         time.Duration
	responseduration time.Duration

	mjpegStream  string
	curImg       bytes.Buffer
	curImgLock   sync.RWMutex
	conChan      chan time.Time
	lastConn     time.Time
	lastConnLock sync.RWMutex
	running      bool
	runningLock  sync.RWMutex
	l            net.Listener
	writer       io.Writer
	handler      http.Handler
}

// NewMjpegproxy returns a new Mjpegproxy with default values.
func NewMjpegproxy() *Mjpegproxy {
	p := &Mjpegproxy{
		// Max MJPEG-frame size 5Mb.
		partbufsize: 625000,
		// Sleep time between error and reconnecting to stream.
		waittime: time.Second * 1,
		// How long to use one stream response before reconnecting.
		responseduration: time.Hour,
	}
	return p
}

// ServeHTTP uses w to serve current last MJPEG-frame
// as JPG. It also reopens MJPEG-stream
// if it was closed by idle timeout.
func (m *Mjpegproxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// No caching for HTTP 1.1.
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	// No caching for HTTP 1.0.
	w.Header().Set("Pragma", "no-cache")
	// No caching for proxies
	w.Header().Set("Expires", "0")
	w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))

	buf := bytes.Buffer{}
	m.curImgLock.RLock()
	buf.Write(m.curImg.Bytes())
	m.curImgLock.RUnlock()
	w.Write(buf.Bytes())

	select {
	case m.conChan <- time.Now():
	default:
		m.lastConnLock.Lock()
		m.lastConn = time.Now()
		m.lastConnLock.Unlock()
	}
}

// CloseStream stops and closes MJPEG-stream.
func (m *Mjpegproxy) CloseStream() {
	m.setRunning(false)
}

// OpenStream creates a go-routine of openstream.
func (m *Mjpegproxy) OpenStream(mjpegStream, user, pass string, timeout time.Duration) {
	go m.openstream(mjpegStream, user, pass, timeout)
}

// GetRunning returns state of openstream.
func (m *Mjpegproxy) GetRunning() bool {
	m.runningLock.RLock()
	defer m.runningLock.RUnlock()
	return m.running
}

func (m *Mjpegproxy) setRunning(r bool) {
	m.runningLock.Lock()
	defer m.runningLock.Unlock()
	m.running = r
}

func (m *Mjpegproxy) getresponse(request *http.Request) (*http.Response, error) {
	tr := &http.Transport{DisableKeepAlives: true}
	client := &http.Client{Transport: tr}
	response, err := client.Do(request)
	if err != nil {
		return nil, err
	}
	if response.StatusCode != 200 {
		response.Body.Close()
		errs := "Got invalid response status: " + response.Status
		return nil, errors.New(errs)
	}
	return response, nil
}

func (m *Mjpegproxy) getboundary(response *http.Response) (string, error) {
	header := response.Header.Get("Content-Type")
	if header == "" {
		return "", errors.New("Content-Type isn't specified!")
	}
	ct, params, err := mime.ParseMediaType(header)
	if err != nil {
		return "", err
	}
	if ct != "multipart/x-mixed-replace" {
		errs := "Wrong Content-Type: expected multipart/x-mixed-replace, got " + ct
		return "", errors.New(errs)
	}
	boundary, ok := params["boundary"]
	if !ok {
		return "", errors.New("No multipart boundary param in Content-Type!")
	}
	// Some IP-cameras screw up boundary strings so we
	// have to remove excessive "--" characters manually.
	boundary = strings.Replace(boundary, "--", "", -1)
	return boundary, nil
}

// OpenStream sends request to target and handles
// response. It opens MJPEG-stream and copies received
// frame to m.curImg. It closes stream if m.CloseStream()
// is called or if difference between current time and
// time of last request to ServeHTTP is bigger than timeout.
func (m *Mjpegproxy) openstream(mjpegStream, user, pass string, timeout time.Duration) {
	m.setRunning(true)
	m.conChan = make(chan time.Time)
	m.mjpegStream = mjpegStream
	var lastconn time.Time
	var img *multipart.Part

	request, err := http.NewRequest("GET", mjpegStream, nil)
	if err != nil {
		log.Fatal(m.mjpegStream, err)
	}
	if user != "" && pass != "" {
		request.SetBasicAuth(user, pass)
	}
	var response *http.Response
	var boundary string
	var mpread *multipart.Reader
	var starttime time.Time

	log.Println("Starting streaming from", mjpegStream)

	for m.GetRunning() {
		lastconn = <-m.conChan
		m.lastConnLock.Lock()
		m.lastConn = lastconn
		m.lastConnLock.Unlock()
		if !m.GetRunning() {
			continue
		}

		response, err = m.getresponse(request)
		if err != nil {
			log.Println(m.mjpegStream, err)
			time.Sleep(m.waittime)
			continue
		}
		starttime = time.Now()
		boundary, err = m.getboundary(response)

		if err != nil {
			log.Println(m.mjpegStream, err)
			response.Body.Close()
			time.Sleep(m.waittime)
			continue
		}
		mpread = multipart.NewReader(response.Body, boundary)
		for m.GetRunning() && (time.Since(lastconn) < timeout) && err == nil {
			if time.Since(starttime) > m.responseduration {
				break
			}
			if time.Since(lastconn) > timeout/2 {
				m.lastConnLock.RLock()
				lastconn = m.lastConn
				m.lastConnLock.RUnlock()
			}
			img, err = mpread.NextPart()
			if err != nil {
				log.Println(m.mjpegStream, err)
				break
			}
			m.curImgLock.Lock()
			m.curImg.Reset()
			_, err = m.curImg.ReadFrom(img)
			m.curImgLock.Unlock()
			img.Close()
			if err != nil {
				log.Println(m.mjpegStream, err)
				break
			}
		}
		response.Body.Close()
		time.Sleep(m.waittime)
	}
	log.Println("Stopped streaming from", mjpegStream)
}
<commit_msg>Add additional buffer<commit_after>// Copyright 2014 Jarmo Puttonen <jarmo.puttonen@gmail.com>. All rights reserved.
// Use of this source code is governed by a MIT-style
// licence that can be found in the LICENCE file.

/*Package paparazzogo implements a caching proxy for
serving MJPEG-stream as JPG-images.
*/
package paparazzogo

import (
	"bytes"
	"errors"
	"io"
	"log"
	"mime"
	"mime/multipart"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// A Mjpegproxy implements http.Handler	interface and generates
// JPG-images from a MJPEG-stream.
type Mjpegproxy struct {
	partbufsize      int64
	waittime         time.Duration
	responseduration time.Duration

	mjpegStream  string
	curImg       bytes.Buffer
	curImgLock   sync.RWMutex
	conChan      chan time.Time
	lastConn     time.Time
	lastConnLock sync.RWMutex
	running      bool
	runningLock  sync.RWMutex
	l            net.Listener
	writer       io.Writer
	handler      http.Handler
}

// NewMjpegproxy returns a new Mjpegproxy with default values.
func NewMjpegproxy() *Mjpegproxy {
	p := &Mjpegproxy{
		// Max MJPEG-frame size 5Mb.
		partbufsize: 625000,
		// Sleep time between error and reconnecting to stream.
		waittime: time.Second * 1,
		// How long to use one stream response before reconnecting.
		responseduration: time.Hour,
	}
	return p
}

// ServeHTTP uses w to serve current last MJPEG-frame
// as JPG. It also reopens MJPEG-stream
// if it was closed by idle timeout.
func (m *Mjpegproxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// No caching for HTTP 1.1.
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	// No caching for HTTP 1.0.
	w.Header().Set("Pragma", "no-cache")
	// No caching for proxies
	w.Header().Set("Expires", "0")
	w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))

	buf := bytes.Buffer{}
	m.curImgLock.RLock()
	buf.Write(m.curImg.Bytes())
	m.curImgLock.RUnlock()
	w.Write(buf.Bytes())

	select {
	case m.conChan <- time.Now():
	default:
		m.lastConnLock.Lock()
		m.lastConn = time.Now()
		m.lastConnLock.Unlock()
	}
}

// CloseStream stops and closes MJPEG-stream.
func (m *Mjpegproxy) CloseStream() {
	m.setRunning(false)
}

// OpenStream creates a go-routine of openstream.
func (m *Mjpegproxy) OpenStream(mjpegStream, user, pass string, timeout time.Duration) {
	go m.openstream(mjpegStream, user, pass, timeout)
}

// GetRunning returns state of openstream.
func (m *Mjpegproxy) GetRunning() bool {
	m.runningLock.RLock()
	defer m.runningLock.RUnlock()
	return m.running
}

func (m *Mjpegproxy) setRunning(r bool) {
	m.runningLock.Lock()
	defer m.runningLock.Unlock()
	m.running = r
}

func (m *Mjpegproxy) getresponse(request *http.Request) (*http.Response, error) {
	tr := &http.Transport{DisableKeepAlives: true}
	client := &http.Client{Transport: tr}
	response, err := client.Do(request)
	if err != nil {
		return nil, err
	}
	if response.StatusCode != 200 {
		response.Body.Close()
		errs := "Got invalid response status: " + response.Status
		return nil, errors.New(errs)
	}
	return response, nil
}

func (m *Mjpegproxy) getboundary(response *http.Response) (string, error) {
	header := response.Header.Get("Content-Type")
	if header == "" {
		return "", errors.New("Content-Type isn't specified!")
	}
	ct, params, err := mime.ParseMediaType(header)
	if err != nil {
		return "", err
	}
	if ct != "multipart/x-mixed-replace" {
		errs := "Wrong Content-Type: expected multipart/x-mixed-replace, got " + ct
		return "", errors.New(errs)
	}
	boundary, ok := params["boundary"]
	if !ok {
		return "", errors.New("No multipart boundary param in Content-Type!")
	}
	// Some IP-cameras screw up boundary strings so we
	// have to remove excessive "--" characters manually.
	boundary = strings.Replace(boundary, "--", "", -1)
	return boundary, nil
}

// OpenStream sends request to target and handles
// response. It opens MJPEG-stream and copies received
// frame to m.curImg. It closes stream if m.CloseStream()
// is called or if difference between current time and
// time of last request to ServeHTTP is bigger than timeout.
func (m *Mjpegproxy) openstream(mjpegStream, user, pass string, timeout time.Duration) {
	m.setRunning(true)
	m.conChan = make(chan time.Time)
	m.mjpegStream = mjpegStream
	var lastconn time.Time
	var img *multipart.Part

	request, err := http.NewRequest("GET", mjpegStream, nil)
	if err != nil {
		log.Fatal(m.mjpegStream, err)
	}
	if user != "" && pass != "" {
		request.SetBasicAuth(user, pass)
	}
	var response *http.Response
	var boundary string
	var mpread *multipart.Reader
	var starttime time.Time
	buf := new(bytes.Buffer)

	log.Println("Starting streaming from", mjpegStream)

	for m.GetRunning() {
		lastconn = <-m.conChan
		m.lastConnLock.Lock()
		m.lastConn = lastconn
		m.lastConnLock.Unlock()
		if !m.GetRunning() {
			continue
		}

		response, err = m.getresponse(request)
		if err != nil {
			log.Println(m.mjpegStream, err)
			time.Sleep(m.waittime)
			continue
		}
		starttime = time.Now()
		boundary, err = m.getboundary(response)

		if err != nil {
			log.Println(m.mjpegStream, err)
			response.Body.Close()
			time.Sleep(m.waittime)
			continue
		}
		mpread = multipart.NewReader(response.Body, boundary)
		for m.GetRunning() && (time.Since(lastconn) < timeout) && err == nil {
			if time.Since(starttime) > m.responseduration {
				break
			}
			if time.Since(lastconn) > timeout/2 {
				m.lastConnLock.RLock()
				lastconn = m.lastConn
				m.lastConnLock.RUnlock()
			}
			img, err = mpread.NextPart()
			if err != nil {
				log.Println(m.mjpegStream, err)
				break
			}
			// buf is an additional buffer that allows
			// serving curImg while loading next part.
			buf.Reset()
			_, err = buf.ReadFrom(img)
			m.curImgLock.Lock()
			m.curImg.Reset()
			_, err = m.curImg.ReadFrom(buf)
			m.curImgLock.Unlock()
			img.Close()
			if err != nil {
				log.Println(m.mjpegStream, err)
				break
			}
		}
		response.Body.Close()
		time.Sleep(m.waittime)
	}
	log.Println("Stopped streaming from", mjpegStream)
}
<|endoftext|>


<commit_before>package digitalocean

import (
	"github.com/mitchellh/packer/packer"
	"testing"
)

func testConfig() map[string]interface{} {
	return map[string]interface{}{
		"client_id": "foo",
		"api_key":   "bar",
	}
}

func TestBuilder_ImplementsBuilder(t *testing.T) {
	var raw interface{}
	raw = &Builder{}
	if _, ok := raw.(packer.Builder); !ok {
		t.Fatalf("Builder should be a builder")
	}
}

func TestBuilder_Prepare_BadType(t *testing.T) {
	b := &Builder{}
	c := map[string]interface{}{
		"api_key": []string{},
	}

	err := b.Prepare(c)
	if err == nil {
		t.Fatalf("prepare should fail")
	}
}

func TestBuilderPrepare_APIKey(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test good
	config["api_key"] = "foo"
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.APIKey != "foo" {
		t.Errorf("access key invalid: %s", b.config.APIKey)
	}

	// Test bad
	delete(config, "api_key")
	b = Builder{}
	err = b.Prepare(config)
	if err == nil {
		t.Fatal("should have error")
	}
}

func TestBuilderPrepare_ClientID(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test good
	config["client_id"] = "foo"
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.ClientID != "foo" {
		t.Errorf("invalid: %s", b.config.ClientID)
	}

	// Test bad
	delete(config, "client_id")
	b = Builder{}
	err = b.Prepare(config)
	if err == nil {
		t.Fatal("should have error")
	}
}

func TestBuilderPrepare_RegionID(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.RegionID != 1 {
		t.Errorf("invalid: %d", b.config.RegionID)
	}

	// Test set
	config["region_id"] = 2
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.RegionID != 2 {
		t.Errorf("invalid: %d", b.config.RegionID)
	}
}

func TestBuilderPrepare_SizeID(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SizeID != 66 {
		t.Errorf("invalid: %d", b.config.SizeID)
	}

	// Test set
	config["size_id"] = 67
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SizeID != 67 {
		t.Errorf("invalid: %d", b.config.SizeID)
	}
}

func TestBuilderPrepare_ImageID(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SizeID != 2676 {
		t.Errorf("invalid: %d", b.config.SizeID)
	}

	// Test set
	config["size_id"] = 2
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SizeID != 2 {
		t.Errorf("invalid: %d", b.config.SizeID)
	}
}

func TestBuilderPrepare_SSHUsername(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SSHUsername != "root" {
		t.Errorf("invalid: %d", b.config.SSHUsername)
	}

	// Test set
	config["ssh_username"] = ""
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SSHPort != 35 {
		t.Errorf("invalid: %d", b.config.SSHPort)
	}
}

func TestBuilderPrepare_SSHTimeout(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.RawSSHTimeout != "1m" {
		t.Errorf("invalid: %d", b.config.RawSSHTimeout)
	}

	// Test set
	config["ssh_timeout"] = "30s"
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	// Test bad
	config["ssh_timeout"] = "tubes"
	b = Builder{}
	err = b.Prepare(config)
	if err == nil {
		t.Fatal("should have error")
	}

}

func TestBuilderPrepare_SnapshotName(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test set
	config["snapshot_name"] = "foo"
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SnapshotName != "foo" {
		t.Errorf("invalid: %s", b.config.SnapshotName)
	}
}
<commit_msg>builder/digitalocean: Make tests pass<commit_after>package digitalocean

import (
	"github.com/mitchellh/packer/packer"
	"testing"
)

func testConfig() map[string]interface{} {
	return map[string]interface{}{
		"client_id": "foo",
		"api_key":   "bar",
	}
}

func TestBuilder_ImplementsBuilder(t *testing.T) {
	var raw interface{}
	raw = &Builder{}
	if _, ok := raw.(packer.Builder); !ok {
		t.Fatalf("Builder should be a builder")
	}
}

func TestBuilder_Prepare_BadType(t *testing.T) {
	b := &Builder{}
	c := map[string]interface{}{
		"api_key": []string{},
	}

	err := b.Prepare(c)
	if err == nil {
		t.Fatalf("prepare should fail")
	}
}

func TestBuilderPrepare_APIKey(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test good
	config["api_key"] = "foo"
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.APIKey != "foo" {
		t.Errorf("access key invalid: %s", b.config.APIKey)
	}

	// Test bad
	delete(config, "api_key")
	b = Builder{}
	err = b.Prepare(config)
	if err == nil {
		t.Fatal("should have error")
	}
}

func TestBuilderPrepare_ClientID(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test good
	config["client_id"] = "foo"
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.ClientID != "foo" {
		t.Errorf("invalid: %s", b.config.ClientID)
	}

	// Test bad
	delete(config, "client_id")
	b = Builder{}
	err = b.Prepare(config)
	if err == nil {
		t.Fatal("should have error")
	}
}

func TestBuilderPrepare_RegionID(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.RegionID != 1 {
		t.Errorf("invalid: %d", b.config.RegionID)
	}

	// Test set
	config["region_id"] = 2
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.RegionID != 2 {
		t.Errorf("invalid: %d", b.config.RegionID)
	}
}

func TestBuilderPrepare_SizeID(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SizeID != 66 {
		t.Errorf("invalid: %d", b.config.SizeID)
	}

	// Test set
	config["size_id"] = 67
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SizeID != 67 {
		t.Errorf("invalid: %d", b.config.SizeID)
	}
}

func TestBuilderPrepare_ImageID(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SizeID != 66 {
		t.Errorf("invalid: %d", b.config.SizeID)
	}

	// Test set
	config["size_id"] = 2
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SizeID != 2 {
		t.Errorf("invalid: %d", b.config.SizeID)
	}
}

func TestBuilderPrepare_SSHUsername(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SSHUsername != "root" {
		t.Errorf("invalid: %d", b.config.SSHUsername)
	}

	// Test set
	config["ssh_username"] = "foo"
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SSHUsername != "foo" {
		t.Errorf("invalid: %s", b.config.SSHUsername)
	}
}

func TestBuilderPrepare_SSHTimeout(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test default
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.RawSSHTimeout != "1m" {
		t.Errorf("invalid: %d", b.config.RawSSHTimeout)
	}

	// Test set
	config["ssh_timeout"] = "30s"
	b = Builder{}
	err = b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	// Test bad
	config["ssh_timeout"] = "tubes"
	b = Builder{}
	err = b.Prepare(config)
	if err == nil {
		t.Fatal("should have error")
	}

}

func TestBuilderPrepare_SnapshotName(t *testing.T) {
	var b Builder
	config := testConfig()

	// Test set
	config["snapshot_name"] = "foo"
	err := b.Prepare(config)
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}

	if b.config.SnapshotName != "foo" {
		t.Errorf("invalid: %s", b.config.SnapshotName)
	}
}
<|endoftext|>


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

package build

import (
	"fmt"
	"os"
	"path"

	"istio.io/release-builder/pkg/model"
	"istio.io/release-builder/pkg/util"
)

// Archive creates the release archive that users will download. This includes the installation templates,
// istioctl, and various tools.
func Archive(manifest model.Manifest) error {
	// First, build all variants of istioctl (linux, osx, windows). gen-charts is required for manifests compiled in to istioctl.
	if err := util.RunMake(manifest, "istio", nil, "gen-charts", "istioctl-all", "istioctl.completion"); err != nil {
		return fmt.Errorf("failed to make istioctl: %v", err)
	}

	// We build archives for each arch. These contain the same thing except arch specific istioctl
	for _, arch := range []string{"linux-amd64", "linux-armv7", "linux-arm64", "osx", "win"} {
		out := path.Join(manifest.Directory, "work", "archive", arch, fmt.Sprintf("istio-%s", manifest.Version))
		if err := os.MkdirAll(out, 0750); err != nil {
			return err
		}

		// Some files we just directly copy into the release archive
		directCopies := []string{
			"LICENSE",
			"README.md",
		}
		for _, file := range directCopies {
			if err := util.CopyFile(path.Join(manifest.RepoDir("istio"), file), path.Join(out, file)); err != nil {
				return err
			}
		}

		// Set up tools/certs. We filter down to only some file patterns
		includePatterns := []string{"README.md", "Makefile*", "common.mk"}
		if err := util.CopyDirFiltered(path.Join(manifest.RepoDir("istio"), "tools", "certs"), path.Join(out, "tools", "certs"), includePatterns); err != nil {
			return err
		}

		// Set up samples. We filter down to only some file patterns
		// TODO - clean this up. We probably include files we don't want and exclude files we do want.
		includePatterns = []string{"*.yaml", "*.md", "*.sh", "*.txt", "*.pem", "*.conf", "*.tpl", "*.json", "Makefile"}
		if err := util.CopyDirFiltered(path.Join(manifest.RepoDir("istio"), "samples"), path.Join(out, "samples"), includePatterns); err != nil {
			return err
		}

		manifestsDir := path.Join(out, "manifests")
		if err := os.MkdirAll(manifestsDir, 0755); err != nil {
			return err
		}
		if err := util.CopyDir(path.Join(manifest.RepoDir("istio"), "manifests", "charts"), manifestsDir); err != nil {
			return err
		}
		if err := util.CopyDir(path.Join(manifest.RepoDir("istio"), "manifests", "examples"), manifestsDir); err != nil {
			return err
		}
		if err := util.CopyDir(path.Join(manifest.RepoDir("istio"), "manifests", "profiles"), manifestsDir); err != nil {
			return err
		}

		if err := sanitizeTemplate(manifest, path.Join(out, "manifests/profiles/default.yaml")); err != nil {
			return fmt.Errorf("failed to sanitize operator charts")
		}
		if err := util.CopyDir(path.Join(manifest.RepoDir("istio"), "operator", "samples"), path.Join(out, "samples/operator")); err != nil {
			return err
		}

		// Write manifest
		if err := writeManifest(manifest, out); err != nil {
			return fmt.Errorf("failed to write manifest: %v", err)
		}

		// Copy the istioctl binary over
		istioctlBinary := fmt.Sprintf("istioctl-%s", arch)
		istioctlDest := "istioctl"
		if arch == "win" {
			istioctlBinary += ".exe"
			istioctlDest += ".exe"
		}
		if err := util.CopyFile(path.Join(manifest.RepoOutDir("istio"), istioctlBinary), path.Join(out, "bin", istioctlDest)); err != nil {
			return err
		}
		if err := os.Chmod(path.Join(out, "bin", istioctlDest), 0755); err != nil {
			return err
		}

		// Copy the istioctl completions files to the tools directory
		completionFiles := []string{"istioctl.bash", "_istioctl"}
		for _, file := range completionFiles {
			if err := util.CopyFile(path.Join(manifest.RepoOutDir("istio"), file), path.Join(out, "tools", file)); err != nil {
				return err
			}
		}

		if err := createArchive(arch, manifest, out); err != nil {
			return err
		}

		if err := createStandaloneIstioctl(arch, manifest, out); err != nil {
			return err
		}
	}
	return nil
}

func createStandaloneIstioctl(arch string, manifest model.Manifest, out string) error {
	var istioctlArchive string
	// Create a stand alone archive for istioctl
	// Windows should use zip, linux and osx tar
	if arch == "win" {
		istioctlArchive = fmt.Sprintf("istioctl-%s-%s.zip", manifest.Version, arch)
		if err := util.ZipFolder(path.Join(out, "bin", "istioctl.exe"), path.Join(out, "bin", istioctlArchive)); err != nil {
			return fmt.Errorf("failed to zip istioctl: %v", err)
		}
	} else {
		istioctlArchive = fmt.Sprintf("istioctl-%s-%s.tar.gz", manifest.Version, arch)
		icmd := util.VerboseCommand("tar", "-czf", istioctlArchive, "istioctl")
		icmd.Dir = path.Join(out, "bin")
		if err := icmd.Run(); err != nil {
			return fmt.Errorf("failed to tar istioctl: %v", err)
		}
	}
	// Copy files over to the output directory
	archivePath := path.Join(out, "bin", istioctlArchive)
	dest := path.Join(manifest.OutDir(), istioctlArchive)
	if err := util.CopyFile(archivePath, dest); err != nil {
		return fmt.Errorf("failed to package %v release archive: %v", arch, err)
	}

	// Create a SHA of the archive
	if err := util.CreateSha(dest); err != nil {
		return fmt.Errorf("failed to package %v: %v", dest, err)
	}
	return nil
}

func createArchive(arch string, manifest model.Manifest, out string) error {
	var archive string
	// Create the archive from all the above files
	// Windows should use zip, linux and osx tar
	if arch == "win" {
		archive = fmt.Sprintf("istio-%s-%s.zip", manifest.Version, arch)
		if err := util.ZipFolder(path.Join(out, "..", fmt.Sprintf("istio-%s", manifest.Version)), path.Join(out, "..", archive)); err != nil {
			return fmt.Errorf("failed to zip istioctl: %v", err)
		}
	} else {
		archive = fmt.Sprintf("istio-%s-%s.tar.gz", manifest.Version, arch)
		cmd := util.VerboseCommand("tar", "-czf", archive, fmt.Sprintf("istio-%s", manifest.Version))
		cmd.Dir = path.Join(out, "..")
		if err := cmd.Run(); err != nil {
			return err
		}
	}

	// Copy files over to the output directory
	archivePath := path.Join(manifest.WorkDir(), "archive", arch, archive)
	dest := path.Join(manifest.OutDir(), archive)
	if err := util.CopyFile(archivePath, dest); err != nil {
		return fmt.Errorf("failed to package %v release archive: %v", arch, err)
	}
	// Create a SHA of the archive
	if err := util.CreateSha(dest); err != nil {
		return fmt.Errorf("failed to package %v: %v", dest, err)
	}
	return nil
}
<commit_msg>Add osx-arm64 target to release (#567)<commit_after>// Copyright Istio Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package build

import (
	"fmt"
	"os"
	"path"

	"istio.io/release-builder/pkg/model"
	"istio.io/release-builder/pkg/util"
)

// Archive creates the release archive that users will download. This includes the installation templates,
// istioctl, and various tools.
func Archive(manifest model.Manifest) error {
	// First, build all variants of istioctl (linux, osx, windows). gen-charts is required for manifests compiled in to istioctl.
	if err := util.RunMake(manifest, "istio", nil, "gen-charts", "istioctl-all", "istioctl.completion"); err != nil {
		return fmt.Errorf("failed to make istioctl: %v", err)
	}

	// We build archives for each arch. These contain the same thing except arch specific istioctl
	for _, arch := range []string{"linux-amd64", "linux-armv7", "linux-arm64", "osx", "osx-arm64", "win"} {
		out := path.Join(manifest.Directory, "work", "archive", arch, fmt.Sprintf("istio-%s", manifest.Version))
		if err := os.MkdirAll(out, 0o750); err != nil {
			return err
		}

		// Some files we just directly copy into the release archive
		directCopies := []string{
			"LICENSE",
			"README.md",
		}
		for _, file := range directCopies {
			if err := util.CopyFile(path.Join(manifest.RepoDir("istio"), file), path.Join(out, file)); err != nil {
				return err
			}
		}

		// Set up tools/certs. We filter down to only some file patterns
		includePatterns := []string{"README.md", "Makefile*", "common.mk"}
		if err := util.CopyDirFiltered(path.Join(manifest.RepoDir("istio"), "tools", "certs"), path.Join(out, "tools", "certs"), includePatterns); err != nil {
			return err
		}

		// Set up samples. We filter down to only some file patterns
		// TODO - clean this up. We probably include files we don't want and exclude files we do want.
		includePatterns = []string{"*.yaml", "*.md", "*.sh", "*.txt", "*.pem", "*.conf", "*.tpl", "*.json", "Makefile"}
		if err := util.CopyDirFiltered(path.Join(manifest.RepoDir("istio"), "samples"), path.Join(out, "samples"), includePatterns); err != nil {
			return err
		}

		manifestsDir := path.Join(out, "manifests")
		if err := os.MkdirAll(manifestsDir, 0o755); err != nil {
			return err
		}
		if err := util.CopyDir(path.Join(manifest.RepoDir("istio"), "manifests", "charts"), manifestsDir); err != nil {
			return err
		}
		if err := util.CopyDir(path.Join(manifest.RepoDir("istio"), "manifests", "examples"), manifestsDir); err != nil {
			return err
		}
		if err := util.CopyDir(path.Join(manifest.RepoDir("istio"), "manifests", "profiles"), manifestsDir); err != nil {
			return err
		}

		if err := sanitizeTemplate(manifest, path.Join(out, "manifests/profiles/default.yaml")); err != nil {
			return fmt.Errorf("failed to sanitize operator charts")
		}
		if err := util.CopyDir(path.Join(manifest.RepoDir("istio"), "operator", "samples"), path.Join(out, "samples/operator")); err != nil {
			return err
		}

		// Write manifest
		if err := writeManifest(manifest, out); err != nil {
			return fmt.Errorf("failed to write manifest: %v", err)
		}

		// Copy the istioctl binary over
		istioctlBinary := fmt.Sprintf("istioctl-%s", arch)
		istioctlDest := "istioctl"
		if arch == "win" {
			istioctlBinary += ".exe"
			istioctlDest += ".exe"
		}
		if err := util.CopyFile(path.Join(manifest.RepoOutDir("istio"), istioctlBinary), path.Join(out, "bin", istioctlDest)); err != nil {
			return err
		}
		if err := os.Chmod(path.Join(out, "bin", istioctlDest), 0o755); err != nil {
			return err
		}

		// Copy the istioctl completions files to the tools directory
		completionFiles := []string{"istioctl.bash", "_istioctl"}
		for _, file := range completionFiles {
			if err := util.CopyFile(path.Join(manifest.RepoOutDir("istio"), file), path.Join(out, "tools", file)); err != nil {
				return err
			}
		}

		if err := createArchive(arch, manifest, out); err != nil {
			return err
		}

		if err := createStandaloneIstioctl(arch, manifest, out); err != nil {
			return err
		}
	}
	return nil
}

func createStandaloneIstioctl(arch string, manifest model.Manifest, out string) error {
	var istioctlArchive string
	// Create a stand alone archive for istioctl
	// Windows should use zip, linux and osx tar
	if arch == "win" {
		istioctlArchive = fmt.Sprintf("istioctl-%s-%s.zip", manifest.Version, arch)
		if err := util.ZipFolder(path.Join(out, "bin", "istioctl.exe"), path.Join(out, "bin", istioctlArchive)); err != nil {
			return fmt.Errorf("failed to zip istioctl: %v", err)
		}
	} else {
		istioctlArchive = fmt.Sprintf("istioctl-%s-%s.tar.gz", manifest.Version, arch)
		icmd := util.VerboseCommand("tar", "-czf", istioctlArchive, "istioctl")
		icmd.Dir = path.Join(out, "bin")
		if err := icmd.Run(); err != nil {
			return fmt.Errorf("failed to tar istioctl: %v", err)
		}
	}
	// Copy files over to the output directory
	archivePath := path.Join(out, "bin", istioctlArchive)
	dest := path.Join(manifest.OutDir(), istioctlArchive)
	if err := util.CopyFile(archivePath, dest); err != nil {
		return fmt.Errorf("failed to package %v release archive: %v", arch, err)
	}

	// Create a SHA of the archive
	if err := util.CreateSha(dest); err != nil {
		return fmt.Errorf("failed to package %v: %v", dest, err)
	}
	return nil
}

func createArchive(arch string, manifest model.Manifest, out string) error {
	var archive string
	// Create the archive from all the above files
	// Windows should use zip, linux and osx tar
	if arch == "win" {
		archive = fmt.Sprintf("istio-%s-%s.zip", manifest.Version, arch)
		if err := util.ZipFolder(path.Join(out, "..", fmt.Sprintf("istio-%s", manifest.Version)), path.Join(out, "..", archive)); err != nil {
			return fmt.Errorf("failed to zip istioctl: %v", err)
		}
	} else {
		archive = fmt.Sprintf("istio-%s-%s.tar.gz", manifest.Version, arch)
		cmd := util.VerboseCommand("tar", "-czf", archive, fmt.Sprintf("istio-%s", manifest.Version))
		cmd.Dir = path.Join(out, "..")
		if err := cmd.Run(); err != nil {
			return err
		}
	}

	// Copy files over to the output directory
	archivePath := path.Join(manifest.WorkDir(), "archive", arch, archive)
	dest := path.Join(manifest.OutDir(), archive)
	if err := util.CopyFile(archivePath, dest); err != nil {
		return fmt.Errorf("failed to package %v release archive: %v", arch, err)
	}
	// Create a SHA of the archive
	if err := util.CreateSha(dest); err != nil {
		return fmt.Errorf("failed to package %v: %v", dest, err)
	}
	return nil
}
<|endoftext|>


<commit_before>package goscp

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/cheggaaa/pb"
	"golang.org/x/crypto/ssh"
)

var (
	// SCP messages
	fileCopyRx  = regexp.MustCompile(`C(?P<mode>\d{4}) (?P<length>\d+) (?P<filename>.+)`)
	dirCopyRx   = regexp.MustCompile(`D(?P<mode>\d{4}) (?P<length>\d+) (?P<dirname>.+)`)
	timestampRx = regexp.MustCompile(`T(?P<mtime>\d+) 0 (?P<atime>\d+) 0`)
	endDir      = "E"
)

type Client struct {
	SSHClient        *ssh.Client
	ProgressCallback func(out string)
	DestinationPath  []string

	// Errors that have occurred while communicating with host
	errors []error

	// Verbose output when communicating with host
	Verbose bool

	// Stop transfer on OS error - occurs during filepath.Walk
	StopOnOSError bool

	// Stdin for SSH session
	scpStdinPipe io.WriteCloser

	// Stdout for SSH session
	scpStdoutPipe *Reader
}

// Returns a ssh.Client wrapper.
// DestinationPath is set to the current directory by default.
func NewClient(c *ssh.Client) *Client {
	return &Client{
		SSHClient:       c,
		DestinationPath: []string{"."},
	}
}

// Set where content will be sent
func (c *Client) SetDestinationPath(path string) {
	c.DestinationPath = []string{path}
}

func (c *Client) addError(err error) {
	c.errors = append(c.errors, err)
}

// GetLastError should be queried after a call to Download() or Upload().
func (c *Client) GetLastError() error {
	if len(c.errors) > 0 {
		return c.errors[len(c.errors)-1]
	}
	return nil
}

// GetErrorStack returns all errors that have occurred so far
func (c *Client) GetErrorStack() []error {
	return c.errors
}

// Cancel an ongoing operation
func (c *Client) Cancel() {
	if c.scpStdoutPipe != nil {
		c.scpStdoutPipe.cancel <- struct{}{}
	}
}

// Download remotePath to c.DestinationPath
func (c *Client) Download(remotePath string) {
	session, err := c.SSHClient.NewSession()
	if err != nil {
		c.addError(err)
		return
	}
	defer session.Close()

	go func() {
		c.scpStdinPipe, err = session.StdinPipe()
		if err != nil {
			c.addError(err)
			return
		}
		defer c.scpStdinPipe.Close()

		r, err := session.StdoutPipe()
		if err != nil {
			c.addError(err)
			return
		}

		// Initialise transfer
		c.sendAck()

		// Wrapper to support cancellation
		c.scpStdoutPipe = &Reader{
			Reader: bufio.NewReader(r),
			cancel: make(chan struct{}, 1),
		}

		for {
			c.outputInfo("Reading message from source")
			msg, err := c.scpStdoutPipe.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					c.addError(err)
				}
				return
			}

			// Strip nulls and new lines
			msg = strings.TrimSpace(strings.Trim(msg, "\x00"))
			c.outputInfo(fmt.Sprintf("Received: %s", msg))

			// Confirm message
			c.sendAck()

			switch {
			case c.isFileCopyMsg(msg):
				// Handle incoming file
				err := c.file(msg)
				if err != nil {
					c.addError(err)
					return
				}
			case c.isDirCopyMsg(msg):
				// Handling incoming directory
				err := c.directory(msg)
				if err != nil {
					c.addError(err)
					return
				}
			case msg == endDir:
				// Directory finished, go up a directory
				c.upDirectory()
			case c.isWarningMsg(msg):
				c.addError(fmt.Errorf("Warning message: [%q]\n", msg))
				return
			case c.isErrorMsg(msg):
				c.addError(fmt.Errorf("Error message: [%q]\n", msg))
				return
			default:
				c.addError(fmt.Errorf("Unhandled message: [%q]\n", msg))
				return
			}

			// Confirm message
			c.sendAck()
		}
	}()

	cmd := fmt.Sprintf("scp -rf %s", remotePath)
	if err := session.Run(cmd); err != nil {
		c.addError(err)
		return
	}

	return
}

// Upload localPath to c.DestinationPath
func (c *Client) Upload(localPath string) {
	session, err := c.SSHClient.NewSession()
	if err != nil {
		c.addError(err)
		return
	}
	defer session.Close()

	go func() {
		c.scpStdinPipe, err = session.StdinPipe()
		if err != nil {
			c.addError(err)
			return
		}
		defer c.scpStdinPipe.Close()

		r, err := session.StdoutPipe()
		if err != nil {
			c.addError(err)
			return
		}

		// Wrapper to support cancellation
		c.scpStdoutPipe = &Reader{
			Reader: bufio.NewReader(r),
			cancel: make(chan struct{}, 1),
		}

		// This has already been used in the cmd call below
		// so it can be reused for 'end of directory' message handling
		c.DestinationPath = []string{}

		err = filepath.Walk(localPath, c.handleItem)
		if err != nil {
			c.addError(err)
			return
		}

		// End transfer
		paths := strings.Split(c.DestinationPath[0], "/")
		for range paths {
			c.sendEndOfDirectoryMessage()
		}
	}()

	cmd := fmt.Sprintf("scp -rt %s", filepath.Join(c.DestinationPath...))
	if err := session.Run(cmd); err != nil {
		c.addError(err)
		return
	}

	return
}

// Send an acknowledgement message
func (c *Client) sendAck() {
	fmt.Fprint(c.scpStdinPipe, "\x00")
}

// Send an error message
func (c *Client) sendErr() {
	fmt.Fprint(c.scpStdinPipe, "\x02")
}

// Check if an incoming message is a file copy message
func (c *Client) isFileCopyMsg(s string) bool {
	return strings.HasPrefix(s, "C")
}

// Check if an incoming message is a directory copy message
func (c *Client) isDirCopyMsg(s string) bool {
	return strings.HasPrefix(s, "D")
}

// Check if an incoming message is a warning
func (c *Client) isWarningMsg(s string) bool {
	return strings.HasPrefix(s, "\x01")
}

// Check if an incoming message is an error
func (c *Client) isErrorMsg(s string) bool {
	return strings.HasPrefix(s, "\x02")
}

// Send a directory message while in source mode
func (c *Client) sendDirectoryMessage(mode os.FileMode, dirname string) {
	msg := fmt.Sprintf("D0%o 0 %s", mode, dirname)
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Send a end of directory message while in source mode
func (c *Client) sendEndOfDirectoryMessage() {
	msg := endDir
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Send a file message while in source mode
func (c *Client) sendFileMessage(mode os.FileMode, size int64, filename string) {
	msg := fmt.Sprintf("C0%o %d %s", mode, size, filename)
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Handle directory copy message in sink mode
func (c *Client) directory(msg string) error {
	parts, err := c.parseMessage(msg, dirCopyRx)
	if err != nil {
		return err
	}

	err = os.Mkdir(filepath.Join(c.DestinationPath...)+string(filepath.Separator)+parts["dirname"], 0755)
	if err != nil {
		return err
	}

	// Traverse into directory
	c.DestinationPath = append(c.DestinationPath, parts["dirname"])

	return nil
}

// Handle file copy message in sink mode
func (c *Client) file(msg string) error {
	parts, err := c.parseMessage(msg, fileCopyRx)
	if err != nil {
		return err
	}

	fileLen, _ := strconv.Atoi(parts["length"])

	// Create local file
	localFile, err := os.Create(filepath.Join(c.DestinationPath...) + string(filepath.Separator) + parts["filename"])
	if err != nil {
		return err
	}
	defer localFile.Close()

	bar := c.newProgressBar(fileLen)
	bar.Start()
	defer bar.Finish()

	mw := io.MultiWriter(localFile, bar)
	if n, err := io.CopyN(mw, c.scpStdoutPipe, int64(fileLen)); err != nil || n < int64(fileLen) {
		c.sendErr()
		return err
	}

	return nil
}

// Break down incoming protocol messages
func (c *Client) parseMessage(msg string, rx *regexp.Regexp) (map[string]string, error) {
	parts := make(map[string]string)
	matches := rx.FindStringSubmatch(msg)
	if len(matches) == 0 {
		return parts, errors.New("Could not parse protocol message: " + msg)
	}

	for i, name := range rx.SubexpNames() {
		parts[name] = matches[i]
	}
	return parts, nil
}

// Go back up one directory
func (c *Client) upDirectory() {
	c.DestinationPath = c.DestinationPath[:len(c.DestinationPath)-1]
}

// Handle each item coming through filepath.Walk
func (c *Client) handleItem(path string, info os.FileInfo, err error) error {
	if err != nil {
		// OS error
		c.outputInfo(fmt.Sprintf("Item error: %s", err))

		if c.StopOnOSError {
			return err
		}
		return nil
	}

	if info.IsDir() {
		// Handle directories
		if len(c.DestinationPath) != 0 {
			// If not first directory
			currentPath := strings.Split(c.DestinationPath[0], "/")
			newPath := strings.Split(path, "/")

			// <= slashes = going back up
			if len(newPath) <= len(currentPath) {
				// Send EOD messages for the amount of directories we go up
				for i := len(newPath) - 1; i < len(currentPath); i++ {
					c.sendEndOfDirectoryMessage()
				}
			}
		}
		c.DestinationPath = []string{path}
		c.sendDirectoryMessage(0644, filepath.Base(path))
	} else {
		// Handle regular files
		targetItem, err := os.Open(path)
		if err != nil {
			return err
		}

		c.sendFileMessage(0644, info.Size(), filepath.Base(path))

		if info.Size() > 0 {
			bar := c.newProgressBar(int(info.Size()))
			bar.Start()
			defer bar.Finish()

			mw := io.MultiWriter(c.scpStdinPipe, bar)

			c.outputInfo(fmt.Sprintf("Sending file: %s", path))
			if _, err := io.Copy(mw, targetItem); err != nil {
				c.sendErr()
				return err
			}

			c.sendAck()
		} else {
			c.outputInfo(fmt.Sprintf("Sending empty file: %s", path))
			c.sendAck()
		}
	}

	return nil
}

func (c *Client) outputInfo(s ...string) {
	if c.Verbose {
		log.Println(s)
	}
}

// Create progress bar
func (c *Client) newProgressBar(fileLength int) *pb.ProgressBar {
	bar := pb.New(fileLength)
	bar.Callback = c.ProgressCallback
	bar.ShowSpeed = true
	bar.ShowTimeLeft = true
	bar.ShowCounters = true
	bar.Units = pb.U_BYTES
	bar.SetRefreshRate(time.Second)
	bar.SetWidth(80)
	bar.SetMaxWidth(80)

	return bar
}

// Wrapper to support cancellation
type Reader struct {
	*bufio.Reader

	// Cancel an ongoing transfer
	cancel chan struct{}
}

// Additional cancellation check
func (r *Reader) Read(p []byte) (n int, err error) {
	select {
	case <-r.cancel:
		return 0, errors.New("Transfer cancelled")
	default:
		return r.Reader.Read(p)
	}
}
<commit_msg>Unexported reader.<commit_after>package goscp

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/cheggaaa/pb"
	"golang.org/x/crypto/ssh"
)

var (
	// SCP messages
	fileCopyRx  = regexp.MustCompile(`C(?P<mode>\d{4}) (?P<length>\d+) (?P<filename>.+)`)
	dirCopyRx   = regexp.MustCompile(`D(?P<mode>\d{4}) (?P<length>\d+) (?P<dirname>.+)`)
	timestampRx = regexp.MustCompile(`T(?P<mtime>\d+) 0 (?P<atime>\d+) 0`)
	endDir      = "E"
)

type Client struct {
	SSHClient        *ssh.Client
	ProgressCallback func(out string)
	DestinationPath  []string

	// Errors that have occurred while communicating with host
	errors []error

	// Verbose output when communicating with host
	Verbose bool

	// Stop transfer on OS error - occurs during filepath.Walk
	StopOnOSError bool

	// Stdin for SSH session
	scpStdinPipe io.WriteCloser

	// Stdout for SSH session
	scpStdoutPipe *reader
}

// Returns a ssh.Client wrapper.
// DestinationPath is set to the current directory by default.
func NewClient(c *ssh.Client) *Client {
	return &Client{
		SSHClient:       c,
		DestinationPath: []string{"."},
	}
}

// Set where content will be sent
func (c *Client) SetDestinationPath(path string) {
	c.DestinationPath = []string{path}
}

func (c *Client) addError(err error) {
	c.errors = append(c.errors, err)
}

// GetLastError should be queried after a call to Download() or Upload().
func (c *Client) GetLastError() error {
	if len(c.errors) > 0 {
		return c.errors[len(c.errors)-1]
	}
	return nil
}

// GetErrorStack returns all errors that have occurred so far
func (c *Client) GetErrorStack() []error {
	return c.errors
}

// Cancel an ongoing operation
func (c *Client) Cancel() {
	if c.scpStdoutPipe != nil {
		c.scpStdoutPipe.cancel <- struct{}{}
	}
}

// Download remotePath to c.DestinationPath
func (c *Client) Download(remotePath string) {
	session, err := c.SSHClient.NewSession()
	if err != nil {
		c.addError(err)
		return
	}
	defer session.Close()

	go func() {
		c.scpStdinPipe, err = session.StdinPipe()
		if err != nil {
			c.addError(err)
			return
		}
		defer c.scpStdinPipe.Close()

		r, err := session.StdoutPipe()
		if err != nil {
			c.addError(err)
			return
		}

		// Initialise transfer
		c.sendAck()

		// Wrapper to support cancellation
		c.scpStdoutPipe = &reader{
			Reader: bufio.NewReader(r),
			cancel: make(chan struct{}, 1),
		}

		for {
			c.outputInfo("Reading message from source")
			msg, err := c.scpStdoutPipe.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					c.addError(err)
				}
				return
			}

			// Strip nulls and new lines
			msg = strings.TrimSpace(strings.Trim(msg, "\x00"))
			c.outputInfo(fmt.Sprintf("Received: %s", msg))

			// Confirm message
			c.sendAck()

			switch {
			case c.isFileCopyMsg(msg):
				// Handle incoming file
				err := c.file(msg)
				if err != nil {
					c.addError(err)
					return
				}
			case c.isDirCopyMsg(msg):
				// Handling incoming directory
				err := c.directory(msg)
				if err != nil {
					c.addError(err)
					return
				}
			case msg == endDir:
				// Directory finished, go up a directory
				c.upDirectory()
			case c.isWarningMsg(msg):
				c.addError(fmt.Errorf("Warning message: [%q]\n", msg))
				return
			case c.isErrorMsg(msg):
				c.addError(fmt.Errorf("Error message: [%q]\n", msg))
				return
			default:
				c.addError(fmt.Errorf("Unhandled message: [%q]\n", msg))
				return
			}

			// Confirm message
			c.sendAck()
		}
	}()

	cmd := fmt.Sprintf("scp -rf %s", remotePath)
	if err := session.Run(cmd); err != nil {
		c.addError(err)
		return
	}

	return
}

// Upload localPath to c.DestinationPath
func (c *Client) Upload(localPath string) {
	session, err := c.SSHClient.NewSession()
	if err != nil {
		c.addError(err)
		return
	}
	defer session.Close()

	go func() {
		c.scpStdinPipe, err = session.StdinPipe()
		if err != nil {
			c.addError(err)
			return
		}
		defer c.scpStdinPipe.Close()

		r, err := session.StdoutPipe()
		if err != nil {
			c.addError(err)
			return
		}

		// Wrapper to support cancellation
		c.scpStdoutPipe = &reader{
			Reader: bufio.NewReader(r),
			cancel: make(chan struct{}, 1),
		}

		// This has already been used in the cmd call below
		// so it can be reused for 'end of directory' message handling
		c.DestinationPath = []string{}

		err = filepath.Walk(localPath, c.handleItem)
		if err != nil {
			c.addError(err)
			return
		}

		// End transfer
		paths := strings.Split(c.DestinationPath[0], "/")
		for range paths {
			c.sendEndOfDirectoryMessage()
		}
	}()

	cmd := fmt.Sprintf("scp -rt %s", filepath.Join(c.DestinationPath...))
	if err := session.Run(cmd); err != nil {
		c.addError(err)
		return
	}

	return
}

// Send an acknowledgement message
func (c *Client) sendAck() {
	fmt.Fprint(c.scpStdinPipe, "\x00")
}

// Send an error message
func (c *Client) sendErr() {
	fmt.Fprint(c.scpStdinPipe, "\x02")
}

// Check if an incoming message is a file copy message
func (c *Client) isFileCopyMsg(s string) bool {
	return strings.HasPrefix(s, "C")
}

// Check if an incoming message is a directory copy message
func (c *Client) isDirCopyMsg(s string) bool {
	return strings.HasPrefix(s, "D")
}

// Check if an incoming message is a warning
func (c *Client) isWarningMsg(s string) bool {
	return strings.HasPrefix(s, "\x01")
}

// Check if an incoming message is an error
func (c *Client) isErrorMsg(s string) bool {
	return strings.HasPrefix(s, "\x02")
}

// Send a directory message while in source mode
func (c *Client) sendDirectoryMessage(mode os.FileMode, dirname string) {
	msg := fmt.Sprintf("D0%o 0 %s", mode, dirname)
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Send a end of directory message while in source mode
func (c *Client) sendEndOfDirectoryMessage() {
	msg := endDir
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Send a file message while in source mode
func (c *Client) sendFileMessage(mode os.FileMode, size int64, filename string) {
	msg := fmt.Sprintf("C0%o %d %s", mode, size, filename)
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Handle directory copy message in sink mode
func (c *Client) directory(msg string) error {
	parts, err := c.parseMessage(msg, dirCopyRx)
	if err != nil {
		return err
	}

	err = os.Mkdir(filepath.Join(c.DestinationPath...)+string(filepath.Separator)+parts["dirname"], 0755)
	if err != nil {
		return err
	}

	// Traverse into directory
	c.DestinationPath = append(c.DestinationPath, parts["dirname"])

	return nil
}

// Handle file copy message in sink mode
func (c *Client) file(msg string) error {
	parts, err := c.parseMessage(msg, fileCopyRx)
	if err != nil {
		return err
	}

	fileLen, _ := strconv.Atoi(parts["length"])

	// Create local file
	localFile, err := os.Create(filepath.Join(c.DestinationPath...) + string(filepath.Separator) + parts["filename"])
	if err != nil {
		return err
	}
	defer localFile.Close()

	bar := c.newProgressBar(fileLen)
	bar.Start()
	defer bar.Finish()

	mw := io.MultiWriter(localFile, bar)
	if n, err := io.CopyN(mw, c.scpStdoutPipe, int64(fileLen)); err != nil || n < int64(fileLen) {
		c.sendErr()
		return err
	}

	return nil
}

// Break down incoming protocol messages
func (c *Client) parseMessage(msg string, rx *regexp.Regexp) (map[string]string, error) {
	parts := make(map[string]string)
	matches := rx.FindStringSubmatch(msg)
	if len(matches) == 0 {
		return parts, errors.New("Could not parse protocol message: " + msg)
	}

	for i, name := range rx.SubexpNames() {
		parts[name] = matches[i]
	}
	return parts, nil
}

// Go back up one directory
func (c *Client) upDirectory() {
	c.DestinationPath = c.DestinationPath[:len(c.DestinationPath)-1]
}

// Handle each item coming through filepath.Walk
func (c *Client) handleItem(path string, info os.FileInfo, err error) error {
	if err != nil {
		// OS error
		c.outputInfo(fmt.Sprintf("Item error: %s", err))

		if c.StopOnOSError {
			return err
		}
		return nil
	}

	if info.IsDir() {
		// Handle directories
		if len(c.DestinationPath) != 0 {
			// If not first directory
			currentPath := strings.Split(c.DestinationPath[0], "/")
			newPath := strings.Split(path, "/")

			// <= slashes = going back up
			if len(newPath) <= len(currentPath) {
				// Send EOD messages for the amount of directories we go up
				for i := len(newPath) - 1; i < len(currentPath); i++ {
					c.sendEndOfDirectoryMessage()
				}
			}
		}
		c.DestinationPath = []string{path}
		c.sendDirectoryMessage(0644, filepath.Base(path))
	} else {
		// Handle regular files
		targetItem, err := os.Open(path)
		if err != nil {
			return err
		}

		c.sendFileMessage(0644, info.Size(), filepath.Base(path))

		if info.Size() > 0 {
			bar := c.newProgressBar(int(info.Size()))
			bar.Start()
			defer bar.Finish()

			mw := io.MultiWriter(c.scpStdinPipe, bar)

			c.outputInfo(fmt.Sprintf("Sending file: %s", path))
			if _, err := io.Copy(mw, targetItem); err != nil {
				c.sendErr()
				return err
			}

			c.sendAck()
		} else {
			c.outputInfo(fmt.Sprintf("Sending empty file: %s", path))
			c.sendAck()
		}
	}

	return nil
}

func (c *Client) outputInfo(s ...string) {
	if c.Verbose {
		log.Println(s)
	}
}

// Create progress bar
func (c *Client) newProgressBar(fileLength int) *pb.ProgressBar {
	bar := pb.New(fileLength)
	bar.Callback = c.ProgressCallback
	bar.ShowSpeed = true
	bar.ShowTimeLeft = true
	bar.ShowCounters = true
	bar.Units = pb.U_BYTES
	bar.SetRefreshRate(time.Second)
	bar.SetWidth(80)
	bar.SetMaxWidth(80)

	return bar
}

// Wrapper to support cancellation
type reader struct {
	*bufio.Reader

	// Cancel an ongoing transfer
	cancel chan struct{}
}

// Additional cancellation check
func (r *reader) Read(p []byte) (n int, err error) {
	select {
	case <-r.cancel:
		return 0, errors.New("Transfer cancelled")
	default:
		return r.Reader.Read(p)
	}
}
<|endoftext|>


<commit_before>// +build arm

package missinggo

import (
	"os"
	"syscall"
	"time"
)

func fileInfoAccessTime(fi os.FileInfo) time.Time {
	ts := fi.Sys().(*syscall.Stat_t).Atim
	return time.Unix(int64(ts.Sec), int64(ts.Nsec))
}
<commit_msg>Remove unnecessary build tag<commit_after>package missinggo

import (
	"os"
	"syscall"
	"time"
)

func fileInfoAccessTime(fi os.FileInfo) time.Time {
	ts := fi.Sys().(*syscall.Stat_t).Atim
	return time.Unix(int64(ts.Sec), int64(ts.Nsec))
}
<|endoftext|>


<commit_before>package model

import (
	"context"
	"fmt"
	"strings"
	"sync/atomic"
	"time"

	"github.com/flimzy/kivik"
	"github.com/flimzy/log"
	"github.com/pkg/errors"

	fb "github.com/FlashbackSRS/flashback-model"
)

func (r *Repo) remoteDSN(name string) string {
	dsn := r.remote.DSN()
	if strings.HasSuffix(dsn, "/") {
		return dsn + name
	}
	return dsn + "/" + name
}

// Sync performs a bi-directional sync.
func (r *Repo) Sync(ctx context.Context) error {
	u, err := r.CurrentUser()
	if err != nil {
		return err
	}
	udbName := "user-" + u
	rdb := r.remoteDSN(udbName)

	var docsWritten, docsRead int32
	if e := r.doSync(ctx, rdb, udbName, &docsWritten, &docsRead); e != nil {
		return errors.Wrap(e, "sync failed")
	}

	updated, err := r.upgradeSchema(ctx)
	if err != nil {
		return errors.Wrap(err, "schema upgrade failed")
	}
	if updated {
		fmt.Printf("Documents were updated\n")
		if e := r.doSync(ctx, rdb, udbName, &docsWritten, &docsRead); e != nil {
			return errors.Wrap(e, "resync failed")
		}
	}

	log.Debugf("Synced %d docs from server, %d to server\n", docsRead, docsWritten)

	return nil
}

func (r *Repo) doSync(ctx context.Context, remoteUserDBName, localUserDBName string, docsWritten, docsRead *int32) error {
	// local to remote
	if e := replicate(ctx, r.local, remoteUserDBName, localUserDBName, docsWritten); e != nil {
		return errors.Wrap(e, "sync local to remote")
	}

	// remote to local
	if e := replicate(ctx, r.local, localUserDBName, remoteUserDBName, docsRead); e != nil {
		return errors.Wrap(e, "sync remote to local")
	}

	if e := r.syncBundles(ctx, docsRead, docsWritten); e != nil {
		return errors.Wrap(e, "bundle sync")
	}

	return errors.Wrap(r.updateSyncTime(ctx), "fialed to store sync timestamp")
}

// upgradeSchema updates the local schema, if necessary, and returns true if
// any updates were made.
//
// This should be run after a sync, and in case of updates, a sync should be
// re-run. This is to reduce the chance of a race condition with multiple
// clients doing a simultaneous update.
func (r *Repo) upgradeSchema(ctx context.Context) (bool, error) {
	defer profile("upgrade")()
	db, err := r.userDB(ctx)
	if err != nil {
		return false, errors.Wrap(err, "failed to connect to db")
	}

	upd := new(int32)
	errs := make(chan error)
	cache := newCardDeckCache(r.local)
	for _, class := range []string{"new", "old", "suspended"} {
		go func(class string) {
			updated, err := upgradeSchemaFromView(ctx, db, cache, class)
			if updated {
				atomic.AddInt32(upd, 1)
			}
			errs <- errors.Wrapf(err, "%s failed", class)
		}(class)
	}
	err = nil
	for i := 0; i < 3; i++ {
		e := <-errs
		if err == nil {
			err = e
		}
	}
	return *upd > 0, err
}

func upgradeSchemaFromView(ctx context.Context, db kivikDB, cache *cardDeckCache, class string) (bool, error) {
	defer profile(fmt.Sprintf("upgrade %s", class))()
	rows, err := db.Query(ctx, "index", "cards", map[string]interface{}{
		"startkey":     []interface{}{class, nil},
		"endkey":       []interface{}{class, nil, map[string]interface{}{}},
		"include_docs": true,
		"reduce":       false,
	})
	if err != nil {
		return false, errors.Wrap(err, "query")
	}
	defer func() { _ = rows.Close() }()
	var count int
	for rows.Next() {
		var card *fb.Card
		if e := rows.ScanDoc(&card); e != nil {
			return count != 0, errors.Wrap(e, "doc scan")
		}
		deckID, err := cache.cardDeck(ctx, card)
		if err != nil {
			return count != 0, errors.Wrap(err, "card deck")
		}
		card.Deck = deckID
		count++
		if _, err := db.Put(ctx, card.ID, card); err != nil {
			return count != 0, errors.Wrap(err, "put")
		}
	}
	log.Debugf("%d of %d %s cards upgraded\n", count, rows.TotalRows(), class)
	return count != 0, errors.Wrap(rows.Err(), "rows")
}

type cardDeckCache struct {
	client      kivikClient
	cache       map[string]string
	readBundles map[string]struct{}
}

func newCardDeckCache(client kivikClient) *cardDeckCache {
	return &cardDeckCache{
		client:      client,
		cache:       make(map[string]string),
		readBundles: make(map[string]struct{}),
	}
}

const orphanedCardDeck = "x"

func (c *cardDeckCache) cardDeck(ctx context.Context, card *fb.Card) (string, error) {
	bundleID := card.BundleID()
	if _, ok := c.readBundles[bundleID]; !ok {
		if err := c.readBundle(ctx, bundleID); err != nil {
			return "", err
		}
	}
	if deckID, ok := c.cache[card.ID]; ok {
		return deckID, nil
	}
	return orphanedCardDeck, nil
}

func (c *cardDeckCache) readBundle(ctx context.Context, bundleID string) error {
	bdb, err := c.client.DB(ctx, bundleID)
	if err != nil {
		return err
	}
	c.readBundles[bundleID] = struct{}{}
	rows, err := bdb.AllDocs(ctx, map[string]interface{}{
		"startkey":     "deck-",
		"endkey":       "deck-" + kivik.EndKeySuffix,
		"include_docs": true,
	})
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var deck fb.Deck
		if err := rows.ScanDoc(&deck); err != nil {
			return err
		}
		for _, cardID := range deck.Cards.All() {
			c.cache[cardID] = deck.ID
		}
	}
	return nil
}

const lastSyncTimestampDocID = "_local/lastSyncTimestamp"

type lastSyncTimestampDoc struct {
	ID       string    `json:"_id"`
	Rev      string    `json:"_rev"`
	LastSync time.Time `json:"lastSync"`
}

// updateSyncTime updates the local timestamp for the last sync.
func (r *Repo) updateSyncTime(ctx context.Context) error {
	rev, _, err := r.lastSyncTime(ctx)
	if err != nil && kivik.StatusCode(err) != kivik.StatusNotFound {
		return err
	}

	u, err := r.CurrentUser()
	if err != nil {
		return err
	}
	db, err := r.local.DB(ctx, "user-"+u)
	if err != nil {
		return err
	}
	doc := lastSyncTimestampDoc{
		ID:       lastSyncTimestampDocID,
		Rev:      rev,
		LastSync: now(),
	}
	_, err = db.Put(ctx, lastSyncTimestampDocID, doc)
	return err
}

// lastSyncTime returns the last time the database was synced.
func (r *Repo) lastSyncTime(ctx context.Context) (rev string, lastSync time.Time, err error) {
	u, err := r.CurrentUser()
	if err != nil {
		return "", time.Time{}, err
	}
	db, err := r.local.DB(ctx, "user-"+u)
	if err != nil {
		return "", time.Time{}, err
	}
	row, err := db.Get(ctx, lastSyncTimestampDocID)
	if err != nil {
		return "", time.Time{}, err
	}
	var doc lastSyncTimestampDoc
	if e := row.ScanDoc(&doc); e != nil {
		return "", time.Time{}, e
	}
	return doc.Rev, doc.LastSync, nil
}

type clientReplicator interface {
	Replicate(context.Context, string, string, ...kivik.Options) (*kivik.Replication, error)
}

func dbDSN(db clientNamer) string {
	dsn := db.Client().DSN()
	dbName := db.Name()
	if dsn != "" && !strings.HasSuffix(dsn, "/") {
		return dsn + "/" + dbName
	}
	return dsn + dbName
}

func replicate(ctx context.Context, client clientReplicator, target, source string, count *int32) error {
	defer profile(fmt.Sprintf("replicate %s -> %s", source, target))()
	replication, err := client.Replicate(ctx, target, source)
	if err != nil {
		return err
	}
	c, err := processReplication(ctx, replication)
	atomic.AddInt32(count, c)
	return err
}

type replication interface {
	IsActive() bool
	Update(context.Context) error
	Delete(context.Context) error
	Err() error
	DocsWritten() int64
}

func processReplication(ctx context.Context, rep replication) (int32, error) {
	// Just wait until the replication is complete
	// TODO: Visual updates
	for rep.IsActive() {
		if err := rep.Update(ctx); err != nil {
			_ = rep.Delete(ctx)
			return int32(rep.DocsWritten()), err
		}
	}
	return int32(rep.DocsWritten()), rep.Err()
}

func (r *Repo) syncBundles(ctx context.Context, reads, writes *int32) error {
	defer profile("syncBundles")()
	udb, err := r.userDB(ctx)
	if err != nil {
		return err
	}
	log.Debugf("Reading bundles from user database...\n")
	rows, err := udb.Find(context.TODO(), map[string]interface{}{
		"selector": map[string]string{"type": "bundle"},
		"fields":   []string{"_id"},
	})
	if err != nil {
		return errors.Wrap(err, "failed to sync bundles")
	}

	var bundles []string
	for rows.Next() {
		var result struct {
			ID string `json:"_id"`
		}
		if err := rows.ScanDoc(&result); err != nil {
			return errors.Wrapf(err, "failed to scan bundle %s", rows.ID())
		}
		bundles = append(bundles, result.ID)
	}
	log.Debugf("bundles = %v\n", bundles)
	for _, bundle := range bundles {
		log.Debugf("Creating remote bundle: %s\n", bundle)
		rdb := r.remoteDSN(bundle)
		if err := r.remote.CreateDB(ctx, bundle); err != nil && kivik.StatusCode(err) != kivik.StatusPreconditionFailed {
			return errors.Wrap(err, "create remote bundle")
		}
		if err := replicate(ctx, r.local, rdb, bundle, writes); err != nil {
			return errors.Wrap(err, "bundle push")
		}
		if err := replicate(ctx, r.local, bundle, rdb, reads); err != nil {
			return errors.Wrap(err, "bundle pull")
		}
	}
	return nil
}

/*
func SyncReviews(local, remote *repo.DB) (int32, error) {
	u, err := repo.CurrentUser()
	if err != nil {
		return 0, err
	}
	host := util.CouchHost()
	ldb, err := util.ReviewsSyncDbs()
	if err != nil {
		return 0, err
	}
	if ldb == nil {
		return 0, nil
	}
	before, err := ldb.Info()
	if err != nil {
		return 0, err
	}
	if before.DocCount == 0 {
		// Nothing at all to sync
		return 0, nil
	}
	rdb, err := repo.NewDB(host + "/" + u.MasterReviewsDBName())
	if err != nil {
		return 0, err
	}
	revsSynced, err := Sync(ldb, rdb)
	if err != nil {
		return 0, err
	}
	after, err := ldb.Info()
	if err != nil {
		return revsSynced, err
	}
	if before.DocCount != after.DocCount || before.UpdateSeq != after.UpdateSeq {
		log.Debugf("ReviewsDb content changed during sync. Refusing to delete.\n")
		return revsSynced, nil
	}
	log.Debugf("Ready to zap %s\n", after.DBName)
	err = util.ZapReviewsDb(ldb)
	return revsSynced, err
}
*/
<commit_msg>Fix shadow variable<commit_after>package model

import (
	"context"
	"fmt"
	"strings"
	"sync/atomic"
	"time"

	"github.com/flimzy/kivik"
	"github.com/flimzy/log"
	"github.com/pkg/errors"

	fb "github.com/FlashbackSRS/flashback-model"
)

func (r *Repo) remoteDSN(name string) string {
	dsn := r.remote.DSN()
	if strings.HasSuffix(dsn, "/") {
		return dsn + name
	}
	return dsn + "/" + name
}

// Sync performs a bi-directional sync.
func (r *Repo) Sync(ctx context.Context) error {
	u, err := r.CurrentUser()
	if err != nil {
		return err
	}
	udbName := "user-" + u
	rdb := r.remoteDSN(udbName)

	var docsWritten, docsRead int32
	if e := r.doSync(ctx, rdb, udbName, &docsWritten, &docsRead); e != nil {
		return errors.Wrap(e, "sync failed")
	}

	updated, err := r.upgradeSchema(ctx)
	if err != nil {
		return errors.Wrap(err, "schema upgrade failed")
	}
	if updated {
		fmt.Printf("Documents were updated\n")
		if e := r.doSync(ctx, rdb, udbName, &docsWritten, &docsRead); e != nil {
			return errors.Wrap(e, "resync failed")
		}
	}

	log.Debugf("Synced %d docs from server, %d to server\n", docsRead, docsWritten)

	return nil
}

func (r *Repo) doSync(ctx context.Context, remoteUserDBName, localUserDBName string, docsWritten, docsRead *int32) error {
	// local to remote
	if e := replicate(ctx, r.local, remoteUserDBName, localUserDBName, docsWritten); e != nil {
		return errors.Wrap(e, "sync local to remote")
	}

	// remote to local
	if e := replicate(ctx, r.local, localUserDBName, remoteUserDBName, docsRead); e != nil {
		return errors.Wrap(e, "sync remote to local")
	}

	if e := r.syncBundles(ctx, docsRead, docsWritten); e != nil {
		return errors.Wrap(e, "bundle sync")
	}

	return errors.Wrap(r.updateSyncTime(ctx), "fialed to store sync timestamp")
}

// upgradeSchema updates the local schema, if necessary, and returns true if
// any updates were made.
//
// This should be run after a sync, and in case of updates, a sync should be
// re-run. This is to reduce the chance of a race condition with multiple
// clients doing a simultaneous update.
func (r *Repo) upgradeSchema(ctx context.Context) (bool, error) {
	defer profile("upgrade")()
	db, err := r.userDB(ctx)
	if err != nil {
		return false, errors.Wrap(err, "failed to connect to db")
	}

	upd := new(int32)
	errs := make(chan error)
	cache := newCardDeckCache(r.local)
	for _, class := range []string{"new", "old", "suspended"} {
		go func(class string) {
			updated, e := upgradeSchemaFromView(ctx, db, cache, class)
			if updated {
				atomic.AddInt32(upd, 1)
			}
			errs <- errors.Wrapf(e, "%s failed", class)
		}(class)
	}
	err = nil
	for i := 0; i < 3; i++ {
		e := <-errs
		if err == nil {
			err = e
		}
	}
	return *upd > 0, err
}

func upgradeSchemaFromView(ctx context.Context, db kivikDB, cache *cardDeckCache, class string) (bool, error) {
	defer profile(fmt.Sprintf("upgrade %s", class))()
	rows, err := db.Query(ctx, "index", "cards", map[string]interface{}{
		"startkey":     []interface{}{class, nil},
		"endkey":       []interface{}{class, nil, map[string]interface{}{}},
		"include_docs": true,
		"reduce":       false,
	})
	if err != nil {
		return false, errors.Wrap(err, "query")
	}
	defer func() { _ = rows.Close() }()
	var count int
	for rows.Next() {
		var card *fb.Card
		if e := rows.ScanDoc(&card); e != nil {
			return count != 0, errors.Wrap(e, "doc scan")
		}
		deckID, err := cache.cardDeck(ctx, card)
		if err != nil {
			return count != 0, errors.Wrap(err, "card deck")
		}
		card.Deck = deckID
		count++
		if _, err := db.Put(ctx, card.ID, card); err != nil {
			return count != 0, errors.Wrap(err, "put")
		}
	}
	log.Debugf("%d of %d %s cards upgraded\n", count, rows.TotalRows(), class)
	return count != 0, errors.Wrap(rows.Err(), "rows")
}

type cardDeckCache struct {
	client      kivikClient
	cache       map[string]string
	readBundles map[string]struct{}
}

func newCardDeckCache(client kivikClient) *cardDeckCache {
	return &cardDeckCache{
		client:      client,
		cache:       make(map[string]string),
		readBundles: make(map[string]struct{}),
	}
}

const orphanedCardDeck = "x"

func (c *cardDeckCache) cardDeck(ctx context.Context, card *fb.Card) (string, error) {
	bundleID := card.BundleID()
	if _, ok := c.readBundles[bundleID]; !ok {
		if err := c.readBundle(ctx, bundleID); err != nil {
			return "", err
		}
	}
	if deckID, ok := c.cache[card.ID]; ok {
		return deckID, nil
	}
	return orphanedCardDeck, nil
}

func (c *cardDeckCache) readBundle(ctx context.Context, bundleID string) error {
	bdb, err := c.client.DB(ctx, bundleID)
	if err != nil {
		return err
	}
	c.readBundles[bundleID] = struct{}{}
	rows, err := bdb.AllDocs(ctx, map[string]interface{}{
		"startkey":     "deck-",
		"endkey":       "deck-" + kivik.EndKeySuffix,
		"include_docs": true,
	})
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var deck fb.Deck
		if err := rows.ScanDoc(&deck); err != nil {
			return err
		}
		for _, cardID := range deck.Cards.All() {
			c.cache[cardID] = deck.ID
		}
	}
	return nil
}

const lastSyncTimestampDocID = "_local/lastSyncTimestamp"

type lastSyncTimestampDoc struct {
	ID       string    `json:"_id"`
	Rev      string    `json:"_rev"`
	LastSync time.Time `json:"lastSync"`
}

// updateSyncTime updates the local timestamp for the last sync.
func (r *Repo) updateSyncTime(ctx context.Context) error {
	rev, _, err := r.lastSyncTime(ctx)
	if err != nil && kivik.StatusCode(err) != kivik.StatusNotFound {
		return err
	}

	u, err := r.CurrentUser()
	if err != nil {
		return err
	}
	db, err := r.local.DB(ctx, "user-"+u)
	if err != nil {
		return err
	}
	doc := lastSyncTimestampDoc{
		ID:       lastSyncTimestampDocID,
		Rev:      rev,
		LastSync: now(),
	}
	_, err = db.Put(ctx, lastSyncTimestampDocID, doc)
	return err
}

// lastSyncTime returns the last time the database was synced.
func (r *Repo) lastSyncTime(ctx context.Context) (rev string, lastSync time.Time, err error) {
	u, err := r.CurrentUser()
	if err != nil {
		return "", time.Time{}, err
	}
	db, err := r.local.DB(ctx, "user-"+u)
	if err != nil {
		return "", time.Time{}, err
	}
	row, err := db.Get(ctx, lastSyncTimestampDocID)
	if err != nil {
		return "", time.Time{}, err
	}
	var doc lastSyncTimestampDoc
	if e := row.ScanDoc(&doc); e != nil {
		return "", time.Time{}, e
	}
	return doc.Rev, doc.LastSync, nil
}

type clientReplicator interface {
	Replicate(context.Context, string, string, ...kivik.Options) (*kivik.Replication, error)
}

func dbDSN(db clientNamer) string {
	dsn := db.Client().DSN()
	dbName := db.Name()
	if dsn != "" && !strings.HasSuffix(dsn, "/") {
		return dsn + "/" + dbName
	}
	return dsn + dbName
}

func replicate(ctx context.Context, client clientReplicator, target, source string, count *int32) error {
	defer profile(fmt.Sprintf("replicate %s -> %s", source, target))()
	replication, err := client.Replicate(ctx, target, source)
	if err != nil {
		return err
	}
	c, err := processReplication(ctx, replication)
	atomic.AddInt32(count, c)
	return err
}

type replication interface {
	IsActive() bool
	Update(context.Context) error
	Delete(context.Context) error
	Err() error
	DocsWritten() int64
}

func processReplication(ctx context.Context, rep replication) (int32, error) {
	// Just wait until the replication is complete
	// TODO: Visual updates
	for rep.IsActive() {
		if err := rep.Update(ctx); err != nil {
			_ = rep.Delete(ctx)
			return int32(rep.DocsWritten()), err
		}
	}
	return int32(rep.DocsWritten()), rep.Err()
}

func (r *Repo) syncBundles(ctx context.Context, reads, writes *int32) error {
	defer profile("syncBundles")()
	udb, err := r.userDB(ctx)
	if err != nil {
		return err
	}
	log.Debugf("Reading bundles from user database...\n")
	rows, err := udb.Find(context.TODO(), map[string]interface{}{
		"selector": map[string]string{"type": "bundle"},
		"fields":   []string{"_id"},
	})
	if err != nil {
		return errors.Wrap(err, "failed to sync bundles")
	}

	var bundles []string
	for rows.Next() {
		var result struct {
			ID string `json:"_id"`
		}
		if err := rows.ScanDoc(&result); err != nil {
			return errors.Wrapf(err, "failed to scan bundle %s", rows.ID())
		}
		bundles = append(bundles, result.ID)
	}
	log.Debugf("bundles = %v\n", bundles)
	for _, bundle := range bundles {
		log.Debugf("Creating remote bundle: %s\n", bundle)
		rdb := r.remoteDSN(bundle)
		if err := r.remote.CreateDB(ctx, bundle); err != nil && kivik.StatusCode(err) != kivik.StatusPreconditionFailed {
			return errors.Wrap(err, "create remote bundle")
		}
		if err := replicate(ctx, r.local, rdb, bundle, writes); err != nil {
			return errors.Wrap(err, "bundle push")
		}
		if err := replicate(ctx, r.local, bundle, rdb, reads); err != nil {
			return errors.Wrap(err, "bundle pull")
		}
	}
	return nil
}

/*
func SyncReviews(local, remote *repo.DB) (int32, error) {
	u, err := repo.CurrentUser()
	if err != nil {
		return 0, err
	}
	host := util.CouchHost()
	ldb, err := util.ReviewsSyncDbs()
	if err != nil {
		return 0, err
	}
	if ldb == nil {
		return 0, nil
	}
	before, err := ldb.Info()
	if err != nil {
		return 0, err
	}
	if before.DocCount == 0 {
		// Nothing at all to sync
		return 0, nil
	}
	rdb, err := repo.NewDB(host + "/" + u.MasterReviewsDBName())
	if err != nil {
		return 0, err
	}
	revsSynced, err := Sync(ldb, rdb)
	if err != nil {
		return 0, err
	}
	after, err := ldb.Info()
	if err != nil {
		return revsSynced, err
	}
	if before.DocCount != after.DocCount || before.UpdateSeq != after.UpdateSeq {
		log.Debugf("ReviewsDb content changed during sync. Refusing to delete.\n")
		return revsSynced, nil
	}
	log.Debugf("Ready to zap %s\n", after.DBName)
	err = util.ZapReviewsDb(ldb)
	return revsSynced, err
}
*/
<|endoftext|>