<commit_before>// Package api is to start the api to be listening on different endpoints
// The API will listen on the port specified in the PulseConfig.toml.
// There are 2 endpoints:
// POST /log/file this will read the file line by line passing in each line to the algorithm
// POST /log/message (in development) this will take a string and pass it directly to the algorithm
package api

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/davecgh/go-spew/spew"

	"github.com/gophergala2016/Pulse/LogPulse/config"
	"github.com/gophergala2016/Pulse/LogPulse/email"
	"github.com/gophergala2016/Pulse/LogPulse/file"
	"github.com/gophergala2016/Pulse/pulse"
)

// Result is used for ResponseWriter in handlers
type Result struct {
	Status  int    `json:"status"`
	Message string `json:"message"`
}

var buffStrings []string
var port int

func init() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Println(r)
			os.Exit(0)
		}
	}()
	val, err := config.Load()
	if err != nil {
		panic(fmt.Errorf("API: %s", err))
	}
	port = val.Port
}

// Start will run the REST API.
func Start() {
	http.HandleFunc("/log/message", StreamLog)
	http.HandleFunc("/log/file", SendFile)

	fmt.Printf("Listening on localhost:%d\n", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port), nil)

}

// StreamLog listens for post request for a string value.
// This string is then passed to the algorithm for analyzing.
func StreamLog(w http.ResponseWriter, r *http.Request) {

	// Checking to see if the request was a post.
	// If not return a 400: bad request
	if r.Method != "POST" {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}

	// Decoding the body of the response.
	// If we could not parse it as json then respond with a 400: bad request
	decoder := json.NewDecoder(r.Body)
	var body struct {
		Message string `json:"message"`
	}
	err := decoder.Decode(&body)
	if err != nil {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}

	//TODO: post the string to the algorthm for analyzing

	// If we were able to decode and send string to algorithm return a 200: success
	w.Header().Set("Content-Type", "application/json")
	result, _ := json.Marshal(Result{200, "success"})
	io.WriteString(w, string(result))

}

// SendFile listens for a POST that has a form field named file and email in the body.
// Using the file field we will download the specified file to the server.
// The email field is used to email the user the results once algorithm is done.
func SendFile(w http.ResponseWriter, r *http.Request) {

	// Checking to see if the request was a post.
	// If not return a 400: bad request
	if r.Method != "POST" {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}
	compressed := false
	// Get the file field from the form in the response.
	// If we cannot parse it a 400 bad request is returned
	f, header, err := r.FormFile("file")
	fmt.Println("Form File")
	if err != nil {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}

	defer f.Close()

	var body struct {
		Email string `json:"email"`
	}
	// Parse the Form in the response.
	// Check if the email field is a valid email.
	// If not return an 400: bad request
	r.ParseForm()
	body.Email = r.Form["email"][0]
	if !email.IsValid(body.Email) {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "email not valid"})
		io.WriteString(w, string(result))
		return
	}
	if err != nil {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}
	fmt.Println("Received bodys", body.Email)
	extension := filepath.Ext(header.Filename)
	filename := header.Filename[0 : len(header.Filename)-len(extension)]

	if extension == ".gz" {
		// Load compressed file on disk
		out, err := os.Create(fmt.Sprintf("%s.gz", filename))
		if err != nil {
			w.Header().Set("Content-Type", "application/json")
			result, _ := json.Marshal(Result{400, "bad request"})
			io.WriteString(w, string(result))
			return
		}

		defer out.Close()

		// Write the content from POST to the file
		_, err = io.Copy(out, f)
		if err != nil {
			w.Header().Set("Content-Type", "application/json")
			result, _ := json.Marshal(Result{400, "gzip copy failed"})
			io.WriteString(w, string(result))
			return
		}

		// Uncompress file
		err = file.UnGZip(fmt.Sprintf("%s.gz", filename))
		if err != nil {
			spew.Dump(err)
			w.Header().Set("Content-Type", "application/json")
			result, _ := json.Marshal(Result{400, "gzip uncompressed failed"})
			io.WriteString(w, string(result))
			return
		}

		compressed = true
	}

	stdIn := make(chan string)
	email.ByPassMail = true // Needs to bypass emails and store in JSON
	email.OutputFile = filename + ".json"
	email.EmailList = []string{body.Email}

	// Run on separat go routine so that we can give users a response on page first.
	go func() {
		start := time.Now()
		// Start the pulse algorithm
		pulse.Run(stdIn, email.SaveToCache)
		line := make(chan string)

		if compressed {
			file.Read(filename, line)
		} else {
			file.StreamRead(f, line)
		}

		for l := range line {
			if l == "EOF" {
				email.ByPassMail = false
				// Once EOF, time to send email from cache JSON storage
				email.DumpBuffer() // Must clear buffer first before sending
				email.SendFromCache(email.OutputFile)
				break
			}
			stdIn <- l
		}
		close(stdIn)

		elapsed := time.Since(start)
		log.Printf("Pulse Algorithm took %s", elapsed)

		// Clean up
		defer func() {
			if compressed {
				err := os.Remove(filename)
				if err != nil {
					fmt.Println("Failed to delete uncompressed file, please delete")
				}

				err = os.Remove(fmt.Sprintf("%s.gz", filename))
				if err != nil {
					fmt.Println("Failed to delete uncompressed file, please delete")
				}
			}

			err = os.Remove(email.OutputFile)
			if err != nil {
				fmt.Println("Failed to delete uncompressed file, please delete")
			}
		}()
	}()

	// Return a 200 success even if algorithm is still going.
	w.Header().Set("Content-Type", "application/json")
	result, _ := json.Marshal(Result{200, "success"})
	io.WriteString(w, string(result))
}
<commit_msg>forgot period<commit_after>// Package api is to start the api to be listening on different endpoints.
// The API will listen on the port specified in the PulseConfig.toml.
// There are 2 endpoints:
// POST /log/file this will read the file line by line passing in each line to the algorithm
// POST /log/message (in development) this will take a string and pass it directly to the algorithm
package api

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/davecgh/go-spew/spew"

	"github.com/gophergala2016/Pulse/LogPulse/config"
	"github.com/gophergala2016/Pulse/LogPulse/email"
	"github.com/gophergala2016/Pulse/LogPulse/file"
	"github.com/gophergala2016/Pulse/pulse"
)

// Result is used for ResponseWriter in handlers
type Result struct {
	Status  int    `json:"status"`
	Message string `json:"message"`
}

var buffStrings []string
var port int

func init() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Println(r)
			os.Exit(0)
		}
	}()
	val, err := config.Load()
	if err != nil {
		panic(fmt.Errorf("API: %s", err))
	}
	port = val.Port
}

// Start will run the REST API.
func Start() {
	http.HandleFunc("/log/message", StreamLog)
	http.HandleFunc("/log/file", SendFile)

	fmt.Printf("Listening on localhost:%d\n", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port), nil)

}

// StreamLog listens for post request for a string value.
// This string is then passed to the algorithm for analyzing.
func StreamLog(w http.ResponseWriter, r *http.Request) {

	// Checking to see if the request was a post.
	// If not return a 400: bad request
	if r.Method != "POST" {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}

	// Decoding the body of the response.
	// If we could not parse it as json then respond with a 400: bad request
	decoder := json.NewDecoder(r.Body)
	var body struct {
		Message string `json:"message"`
	}
	err := decoder.Decode(&body)
	if err != nil {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}

	//TODO: post the string to the algorthm for analyzing

	// If we were able to decode and send string to algorithm return a 200: success
	w.Header().Set("Content-Type", "application/json")
	result, _ := json.Marshal(Result{200, "success"})
	io.WriteString(w, string(result))

}

// SendFile listens for a POST that has a form field named file and email in the body.
// Using the file field we will download the specified file to the server.
// The email field is used to email the user the results once algorithm is done.
func SendFile(w http.ResponseWriter, r *http.Request) {

	// Checking to see if the request was a post.
	// If not return a 400: bad request
	if r.Method != "POST" {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}
	compressed := false
	// Get the file field from the form in the response.
	// If we cannot parse it a 400 bad request is returned
	f, header, err := r.FormFile("file")
	fmt.Println("Form File")
	if err != nil {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}

	defer f.Close()

	var body struct {
		Email string `json:"email"`
	}
	// Parse the Form in the response.
	// Check if the email field is a valid email.
	// If not return an 400: bad request
	r.ParseForm()
	body.Email = r.Form["email"][0]
	if !email.IsValid(body.Email) {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "email not valid"})
		io.WriteString(w, string(result))
		return
	}
	if err != nil {
		w.Header().Set("Content-Type", "application/json")
		result, _ := json.Marshal(Result{400, "bad request"})
		io.WriteString(w, string(result))
		return
	}
	fmt.Println("Received bodys", body.Email)
	extension := filepath.Ext(header.Filename)
	filename := header.Filename[0 : len(header.Filename)-len(extension)]

	if extension == ".gz" {
		// Load compressed file on disk
		out, err := os.Create(fmt.Sprintf("%s.gz", filename))
		if err != nil {
			w.Header().Set("Content-Type", "application/json")
			result, _ := json.Marshal(Result{400, "bad request"})
			io.WriteString(w, string(result))
			return
		}

		defer out.Close()

		// Write the content from POST to the file
		_, err = io.Copy(out, f)
		if err != nil {
			w.Header().Set("Content-Type", "application/json")
			result, _ := json.Marshal(Result{400, "gzip copy failed"})
			io.WriteString(w, string(result))
			return
		}

		// Uncompress file
		err = file.UnGZip(fmt.Sprintf("%s.gz", filename))
		if err != nil {
			spew.Dump(err)
			w.Header().Set("Content-Type", "application/json")
			result, _ := json.Marshal(Result{400, "gzip uncompressed failed"})
			io.WriteString(w, string(result))
			return
		}

		compressed = true
	}

	stdIn := make(chan string)
	email.ByPassMail = true // Needs to bypass emails and store in JSON
	email.OutputFile = filename + ".json"
	email.EmailList = []string{body.Email}

	// Run on separat go routine so that we can give users a response on page first.
	go func() {
		start := time.Now()
		// Start the pulse algorithm
		pulse.Run(stdIn, email.SaveToCache)
		line := make(chan string)

		if compressed {
			file.Read(filename, line)
		} else {
			file.StreamRead(f, line)
		}

		for l := range line {
			if l == "EOF" {
				email.ByPassMail = false
				// Once EOF, time to send email from cache JSON storage
				email.DumpBuffer() // Must clear buffer first before sending
				email.SendFromCache(email.OutputFile)
				break
			}
			stdIn <- l
		}
		close(stdIn)

		elapsed := time.Since(start)
		log.Printf("Pulse Algorithm took %s", elapsed)

		// Clean up
		defer func() {
			if compressed {
				err := os.Remove(filename)
				if err != nil {
					fmt.Println("Failed to delete uncompressed file, please delete")
				}

				err = os.Remove(fmt.Sprintf("%s.gz", filename))
				if err != nil {
					fmt.Println("Failed to delete uncompressed file, please delete")
				}
			}

			err = os.Remove(email.OutputFile)
			if err != nil {
				fmt.Println("Failed to delete uncompressed file, please delete")
			}
		}()
	}()

	// Return a 200 success even if algorithm is still going.
	w.Header().Set("Content-Type", "application/json")
	result, _ := json.Marshal(Result{200, "success"})
	io.WriteString(w, string(result))
}
<|endoftext|>


<commit_before>package net

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"testing"
	"time"

	"github.com/sahib/brig/backend"
	"github.com/sahib/brig/catfs"
	ie "github.com/sahib/brig/catfs/errors"
	"github.com/sahib/brig/net/peer"
	"github.com/sahib/brig/repo"
	"github.com/stretchr/testify/require"
)

type testUnit struct {
	ctl *Client
	srv *Server
	fs  *catfs.FS
	rp  *repo.Repository
	bk  backend.Backend
}

func withNetServer(t *testing.T, name string, backendPort int, basePath string, fn func(u testUnit)) {
	basePath, err := ioutil.TempDir("", "brig-ctl-test")
	require.Nil(t, err)

	netDbPath := "/tmp/brig-net-test-netdb"
	regDbPath := "/tmp/brig-net-test-reg.yml"
	require.Nil(t, err)

	defer func() {
		require.Nil(t, os.RemoveAll(basePath))
		require.Nil(t, os.RemoveAll(netDbPath))
		require.Nil(t, os.RemoveAll(regDbPath))
	}()

	// The following env vars are only read in FromName.
	require.Nil(t, os.Setenv("BRIG_MOCK_PORT", fmt.Sprintf("%d", backendPort)))
	require.Nil(t, os.Setenv("BRIG_MOCK_USER", name))
	require.Nil(t, os.Setenv("BRIG_MOCK_NET_DB_PATH", netDbPath))
	require.Nil(t, os.Setenv("BRIG_REGISTRY_PATH", regDbPath))
	bk, err := backend.FromName("mock", basePath, nil)
	require.Nil(t, err)

	err = repo.Init(basePath, name, "password", "mock", 6666)
	require.Nil(t, err)

	rp, err := repo.Open(basePath, "password")
	require.Nil(t, err)

	srv, err := NewServer(rp, bk)
	require.Nil(t, err)

	fs, err := rp.FS(name, bk)
	require.Nil(t, err)

	waitForDeath := make(chan bool)
	go func() {
		defer func() {
			waitForDeath <- true
		}()
		require.Nil(t, srv.Serve())
		require.Nil(t, srv.Close())
	}()

	// Allow a short time for the server go routine to fully boot up.
	time.Sleep(50 * time.Millisecond)

	// Run the actual test function:
	fn(testUnit{
		rp:  rp,
		bk:  bk,
		fs:  fs,
		srv: srv,
	})

	// wait until serve was done.
	srv.Quit()
	<-waitForDeath
}

func buildFingerprint(t *testing.T, u testUnit) peer.Fingerprint {
	ownPubKey, err := u.rp.Keyring().OwnPubKey()
	require.Nil(t, err)

	self, err := u.srv.Identity()
	require.Nil(t, err)

	return peer.BuildFingerprint(self.Addr, ownPubKey)
}

func withNetPair(t *testing.T, fn func(a, b testUnit)) {
	basePath, err := ioutil.TempDir("", "brig-net-test")
	require.Nil(t, err)

	defer func() {
		require.Nil(t, os.RemoveAll(basePath))
	}()

	withNetServer(t, "alice", 9998, basePath, func(a testUnit) {
		withNetServer(t, "bob", 9999, basePath, func(b testUnit) {
			// Add each other's fingerprints:
			require.Nil(t, a.rp.Remotes.AddOrUpdateRemote(repo.Remote{
				Name:        "bob",
				Fingerprint: buildFingerprint(t, b),
			}))
			require.Nil(t, b.rp.Remotes.AddOrUpdateRemote(repo.Remote{
				Name:        "alice",
				Fingerprint: buildFingerprint(t, a),
			}))

			ctx := context.Background()
			aliCtl, err := Dial(ctx, "alice", b.rp, b.bk)
			require.Nil(t, err)

			bobCtl, err := Dial(ctx, "bob", a.rp, a.bk)
			require.Nil(t, err)

			a.ctl = bobCtl
			b.ctl = aliCtl

			fn(a, b)
		})
	})
}

func TestClientPing(t *testing.T) {
	withNetPair(t, func(a, b testUnit) {
		for i := 0; i < 100; i++ {
			if err := a.ctl.Ping(); err != nil {
				t.Fatalf("ping to bob failed: %v", err)
			}

			if err := b.ctl.Ping(); err != nil {
				t.Fatalf("ping to alice failed: %v", err)
			}
		}
	})
}

func TestClientFetchStore(t *testing.T) {
	withNetPair(t, func(a, b testUnit) {
		filePath := "/a/new/name/has/been/born"
		fileData := []byte{1, 2, 3}
		fileSrc := bytes.NewReader(fileData)

		if err := a.fs.Stage(filePath, fileSrc); err != nil {
			t.Fatalf("failed to stage simple file: %v", err)
		}

		data, err := b.ctl.FetchStore()
		if err != nil {
			t.Fatalf("failed to read store: %v", err)
		}

		aliceFs, err := b.rp.FS("alice", b.bk)
		if err != nil {
			t.Fatalf("Failed to get empty bob fs: %v", err)
		}

		_, err = aliceFs.Stat(filePath)
		if !ie.IsNoSuchFileError(err) {
			t.Fatalf("File has existed in bob's empty store (wtf?): %v", err)
		}

		if err := aliceFs.Import(data); err != nil {
			t.Fatalf("Failed to import data: %v", err)
		}

		info, err := aliceFs.Stat(filePath)
		if err != nil {
			t.Fatalf("Failed to read file exported from alice: %v", err)
		}

		// Check superficially that store was imported right:
		require.Equal(t, info.Path, filePath)
		require.Equal(t, info.User, "alice")
		require.Equal(t, info.Size, uint64(3))
		require.Equal(t, info.IsDir, false)
	})
}

func TestClientFetchPatch(t *testing.T) {
	withNetPair(t, func(a, b testUnit) {
		// Create a new file in alice's fs.
		require.Nil(t, a.fs.Stage("/new_file", bytes.NewReader([]byte{1, 2, 3})))

		// Get a patch from alice:
		patchData, err := b.ctl.FetchPatch(0)
		require.Nil(t, err)
		require.NotNil(t, patchData)

		// Create a new empty FS for alice.
		aliceFs, err := b.rp.FS("alice", b.bk)
		if err != nil {
			t.Fatalf("Failed to get empty bob fs: %v", err)
		}

		// It should have the initial patch version of 0.
		lastPatchIdx, err := aliceFs.LastPatchIndex()
		require.Nil(t, err)
		require.Equal(t, int64(0), lastPatchIdx)

		// After applying the patch, we should have bob's data.
		require.Nil(t, aliceFs.ApplyPatch(patchData))
		newFileInfo, err := aliceFs.Stat("/new_file")
		require.Nil(t, err)
		require.Equal(t, "/new_file", newFileInfo.Path)
		require.Equal(t, uint64(3), newFileInfo.Size)

		// Check that the new patch version is 2 (i.e. bob has two commits)
		lastPatchIdx, err = aliceFs.LastPatchIndex()
		require.Nil(t, err)
		require.Equal(t, int64(2), lastPatchIdx)

		// If we fetch the same patch again, it will be empty.
		// (data will be not len=0, but no real contents)
		patchData, err = b.ctl.FetchPatch(2)
		require.Nil(t, err)
		require.NotNil(t, patchData)
		require.Nil(t, aliceFs.ApplyPatch(patchData))

		// Last patch was empty, so should not bump the version.
		lastPatchIdx, err = aliceFs.LastPatchIndex()
		require.Nil(t, err)
		require.Equal(t, int64(2), lastPatchIdx)

		// Bob's patch index should not have changed.
		// For bob, the patch index does not make really sense,
		// since he's the owner of the fs and always has the latest version.
		lastBobPatchIdx, err := b.fs.LastPatchIndex()
		require.Nil(t, err)
		require.Equal(t, int64(0), lastBobPatchIdx)
	})
}

func TestClientCompleteFetchAllowed(t *testing.T) {
	withNetPair(t, func(a, b testUnit) {
		isAllowed, err := b.ctl.IsCompleteFetchAllowed()
		require.Nil(t, err)
		require.True(t, isAllowed)

		// Make the remote have only access to a specific sub folder:
		rmt, err := a.rp.Remotes.Remote("bob")
		require.Nil(t, err)

		err = a.rp.Remotes.AddOrUpdateRemote("bob", repo.Remote{
			Fingerprint: rmt.Fingerprint,
			Name:        rmt.Name,
			Folders: []repo.Folder{
				{
					Folder: "/photos",
				},
			},
		})
		require.Nil(t, err)

		isAllowed, err = b.ctl.IsCompleteFetchAllowed()
		require.Nil(t, err)
		require.False(t, isAllowed)

		// Try again with the root folder enabled:
		err = a.rp.Remotes.AddOrUpdateRemote("bob", repo.Remote{
			Fingerprint: rmt.Fingerprint,
			Name:        rmt.Name,
			Folders: []repo.Folder{
				{
					Folder: "/",
				},
				{
					Folder: "/photos",
				},
			},
		})
		require.Nil(t, err)

		// This also tests that we are able to change remote config
		// without needing to reconnect.
		isAllowed, err = b.ctl.IsCompleteFetchAllowed()
		require.Nil(t, err)
		require.True(t, isAllowed)
	})
}
<commit_msg>net: fix test compile error<commit_after>package net

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"testing"
	"time"

	"github.com/sahib/brig/backend"
	"github.com/sahib/brig/catfs"
	ie "github.com/sahib/brig/catfs/errors"
	"github.com/sahib/brig/net/peer"
	"github.com/sahib/brig/repo"
	"github.com/stretchr/testify/require"
)

type testUnit struct {
	ctl *Client
	srv *Server
	fs  *catfs.FS
	rp  *repo.Repository
	bk  backend.Backend
}

func withNetServer(t *testing.T, name string, backendPort int, basePath string, fn func(u testUnit)) {
	basePath, err := ioutil.TempDir("", "brig-ctl-test")
	require.Nil(t, err)

	netDbPath := "/tmp/brig-net-test-netdb"
	regDbPath := "/tmp/brig-net-test-reg.yml"
	require.Nil(t, err)

	defer func() {
		require.Nil(t, os.RemoveAll(basePath))
		require.Nil(t, os.RemoveAll(netDbPath))
		require.Nil(t, os.RemoveAll(regDbPath))
	}()

	// The following env vars are only read in FromName.
	require.Nil(t, os.Setenv("BRIG_MOCK_PORT", fmt.Sprintf("%d", backendPort)))
	require.Nil(t, os.Setenv("BRIG_MOCK_USER", name))
	require.Nil(t, os.Setenv("BRIG_MOCK_NET_DB_PATH", netDbPath))
	require.Nil(t, os.Setenv("BRIG_REGISTRY_PATH", regDbPath))
	bk, err := backend.FromName("mock", basePath, nil)
	require.Nil(t, err)

	err = repo.Init(basePath, name, "password", "mock", 6666)
	require.Nil(t, err)

	rp, err := repo.Open(basePath, "password")
	require.Nil(t, err)

	srv, err := NewServer(rp, bk)
	require.Nil(t, err)

	fs, err := rp.FS(name, bk)
	require.Nil(t, err)

	waitForDeath := make(chan bool)
	go func() {
		defer func() {
			waitForDeath <- true
		}()
		require.Nil(t, srv.Serve())
		require.Nil(t, srv.Close())
	}()

	// Allow a short time for the server go routine to fully boot up.
	time.Sleep(50 * time.Millisecond)

	// Run the actual test function:
	fn(testUnit{
		rp:  rp,
		bk:  bk,
		fs:  fs,
		srv: srv,
	})

	// wait until serve was done.
	srv.Quit()
	<-waitForDeath
}

func buildFingerprint(t *testing.T, u testUnit) peer.Fingerprint {
	ownPubKey, err := u.rp.Keyring().OwnPubKey()
	require.Nil(t, err)

	self, err := u.srv.Identity()
	require.Nil(t, err)

	return peer.BuildFingerprint(self.Addr, ownPubKey)
}

func withNetPair(t *testing.T, fn func(a, b testUnit)) {
	basePath, err := ioutil.TempDir("", "brig-net-test")
	require.Nil(t, err)

	defer func() {
		require.Nil(t, os.RemoveAll(basePath))
	}()

	withNetServer(t, "alice", 9998, basePath, func(a testUnit) {
		withNetServer(t, "bob", 9999, basePath, func(b testUnit) {
			// Add each other's fingerprints:
			require.Nil(t, a.rp.Remotes.AddOrUpdateRemote(repo.Remote{
				Name:        "bob",
				Fingerprint: buildFingerprint(t, b),
			}))
			require.Nil(t, b.rp.Remotes.AddOrUpdateRemote(repo.Remote{
				Name:        "alice",
				Fingerprint: buildFingerprint(t, a),
			}))

			ctx := context.Background()
			aliCtl, err := Dial(ctx, "alice", b.rp, b.bk)
			require.Nil(t, err)

			bobCtl, err := Dial(ctx, "bob", a.rp, a.bk)
			require.Nil(t, err)

			a.ctl = bobCtl
			b.ctl = aliCtl

			fn(a, b)
		})
	})
}

func TestClientPing(t *testing.T) {
	withNetPair(t, func(a, b testUnit) {
		for i := 0; i < 100; i++ {
			if err := a.ctl.Ping(); err != nil {
				t.Fatalf("ping to bob failed: %v", err)
			}

			if err := b.ctl.Ping(); err != nil {
				t.Fatalf("ping to alice failed: %v", err)
			}
		}
	})
}

func TestClientFetchStore(t *testing.T) {
	withNetPair(t, func(a, b testUnit) {
		filePath := "/a/new/name/has/been/born"
		fileData := []byte{1, 2, 3}
		fileSrc := bytes.NewReader(fileData)

		if err := a.fs.Stage(filePath, fileSrc); err != nil {
			t.Fatalf("failed to stage simple file: %v", err)
		}

		data, err := b.ctl.FetchStore()
		if err != nil {
			t.Fatalf("failed to read store: %v", err)
		}

		aliceFs, err := b.rp.FS("alice", b.bk)
		if err != nil {
			t.Fatalf("Failed to get empty bob fs: %v", err)
		}

		_, err = aliceFs.Stat(filePath)
		if !ie.IsNoSuchFileError(err) {
			t.Fatalf("File has existed in bob's empty store (wtf?): %v", err)
		}

		if err := aliceFs.Import(data); err != nil {
			t.Fatalf("Failed to import data: %v", err)
		}

		info, err := aliceFs.Stat(filePath)
		if err != nil {
			t.Fatalf("Failed to read file exported from alice: %v", err)
		}

		// Check superficially that store was imported right:
		require.Equal(t, info.Path, filePath)
		require.Equal(t, info.User, "alice")
		require.Equal(t, info.Size, uint64(3))
		require.Equal(t, info.IsDir, false)
	})
}

func TestClientFetchPatch(t *testing.T) {
	withNetPair(t, func(a, b testUnit) {
		// Create a new file in alice's fs.
		require.Nil(t, a.fs.Stage("/new_file", bytes.NewReader([]byte{1, 2, 3})))

		// Get a patch from alice:
		patchData, err := b.ctl.FetchPatch(0)
		require.Nil(t, err)
		require.NotNil(t, patchData)

		// Create a new empty FS for alice.
		aliceFs, err := b.rp.FS("alice", b.bk)
		if err != nil {
			t.Fatalf("Failed to get empty bob fs: %v", err)
		}

		// It should have the initial patch version of 0.
		lastPatchIdx, err := aliceFs.LastPatchIndex()
		require.Nil(t, err)
		require.Equal(t, int64(0), lastPatchIdx)

		// After applying the patch, we should have bob's data.
		require.Nil(t, aliceFs.ApplyPatch(patchData))
		newFileInfo, err := aliceFs.Stat("/new_file")
		require.Nil(t, err)
		require.Equal(t, "/new_file", newFileInfo.Path)
		require.Equal(t, uint64(3), newFileInfo.Size)

		// Check that the new patch version is 2 (i.e. bob has two commits)
		lastPatchIdx, err = aliceFs.LastPatchIndex()
		require.Nil(t, err)
		require.Equal(t, int64(2), lastPatchIdx)

		// If we fetch the same patch again, it will be empty.
		// (data will be not len=0, but no real contents)
		patchData, err = b.ctl.FetchPatch(2)
		require.Nil(t, err)
		require.NotNil(t, patchData)
		require.Nil(t, aliceFs.ApplyPatch(patchData))

		// Last patch was empty, so should not bump the version.
		lastPatchIdx, err = aliceFs.LastPatchIndex()
		require.Nil(t, err)
		require.Equal(t, int64(2), lastPatchIdx)

		// Bob's patch index should not have changed.
		// For bob, the patch index does not make really sense,
		// since he's the owner of the fs and always has the latest version.
		lastBobPatchIdx, err := b.fs.LastPatchIndex()
		require.Nil(t, err)
		require.Equal(t, int64(0), lastBobPatchIdx)
	})
}

func TestClientCompleteFetchAllowed(t *testing.T) {
	withNetPair(t, func(a, b testUnit) {
		isAllowed, err := b.ctl.IsCompleteFetchAllowed()
		require.Nil(t, err)
		require.True(t, isAllowed)

		// Make the remote have only access to a specific sub folder:
		rmt, err := a.rp.Remotes.Remote("bob")
		require.Nil(t, err)

		err = a.rp.Remotes.AddOrUpdateRemote(repo.Remote{
			Fingerprint: rmt.Fingerprint,
			Name:        rmt.Name,
			Folders: []repo.Folder{
				{
					Folder: "/photos",
				},
			},
		})
		require.Nil(t, err)

		isAllowed, err = b.ctl.IsCompleteFetchAllowed()
		require.Nil(t, err)
		require.False(t, isAllowed)

		// Try again with the root folder enabled:
		err = a.rp.Remotes.AddOrUpdateRemote(repo.Remote{
			Fingerprint: rmt.Fingerprint,
			Name:        rmt.Name,
			Folders: []repo.Folder{
				{
					Folder: "/",
				},
				{
					Folder: "/photos",
				},
			},
		})
		require.Nil(t, err)

		// This also tests that we are able to change remote config
		// without needing to reconnect.
		isAllowed, err = b.ctl.IsCompleteFetchAllowed()
		require.Nil(t, err)
		require.True(t, isAllowed)
	})
}
<|endoftext|>


<commit_before>package monitor

type Target struct {
	Title string
	URL   string
}

type TargetsGetter interface {
	GetTargets() ([]Target, error)
}
<commit_msg>:hammer: Add Target.ID field<commit_after>package monitor

type Target struct {
	ID    uint
	Title string
	URL   string
}

type TargetsGetter interface {
	GetTargets() ([]Target, error)
}
<|endoftext|>


<commit_before>package monitor

import (
	"expvar"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"sort"
	"strconv"
	"sync"
	"time"

	"github.com/influxdb/influxdb/influxql"
)

// Client is the interface modules must implement if they wish to register with monitor.
type Client interface {
	Statistics() (map[string]interface{}, error)
	Diagnostics() (map[string]interface{}, error)
}

// Service represents an instance of the monitor service.
type Service struct {
	wg            sync.WaitGroup
	done          chan struct{}
	mu            sync.Mutex
	registrations []*clientWithMeta

	hostname  string
	clusterID uint64
	nodeID    uint64

	storeEnabled  bool
	storeDatabase string
	storeAddress  string
	storeInterval time.Duration

	Logger *log.Logger
}

// NewService returns a new instance of the monitor service.
func NewService(c Config) *Service {
	return &Service{
		registrations: make([]*clientWithMeta, 0),
		storeEnabled:  c.StoreEnabled,
		storeDatabase: c.StoreDatabase,
		storeAddress:  c.StoreAddress,
		storeInterval: time.Duration(c.StoreInterval),
		Logger:        log.New(os.Stderr, "[monitor] ", log.LstdFlags),
	}
}

// Open opens the monitoring service, using the given clusterID, node ID, and hostname
// for identification purposes.
func (s *Service) Open(clusterID, nodeID uint64, hostname string) error {
	s.Logger.Printf("starting monitor service for cluster %d, host %s", clusterID, hostname)
	s.clusterID = clusterID
	s.nodeID = nodeID
	s.hostname = hostname

	// Self-register Go runtime stats.
	s.Register("runtime", nil, &goRuntime{})

	// If enabled, record stats in a InfluxDB system.
	if s.storeEnabled {
		s.Logger.Printf("storing in %s, database '%s', interval %s",
			s.storeAddress, s.storeDatabase, s.storeInterval)

		s.Logger.Printf("ensuring database %s exists on %s", s.storeDatabase, s.storeAddress)
		if err := ensureDatabaseExists(s.storeAddress, s.storeDatabase); err != nil {
			return err
		}

		// Start periodic writes to system.
		s.wg.Add(1)
		go s.storeStatistics()
	}

	return nil
}

// Close closes the monitor service.
func (s *Service) Close() {
	s.Logger.Println("shutting down monitor service")
	close(s.done)
	s.wg.Wait()
	s.done = nil
}

// SetLogger sets the internal logger to the logger passed in.
func (s *Service) SetLogger(l *log.Logger) {
	s.Logger = l
}

// Register registers a client with the given name and tags.
func (s *Service) Register(name string, tags map[string]string, client Client) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	c := &clientWithMeta{
		Client: client,
		name:   name,
		tags:   tags,
	}
	s.registrations = append(s.registrations, c)
	s.Logger.Printf(`'%s:%v' registered for monitoring`, name, tags)
	return nil
}

// ExecuteStatement executes monitor-related query statements.
func (s *Service) ExecuteStatement(stmt influxql.Statement) *influxql.Result {
	switch stmt := stmt.(type) {
	case *influxql.ShowStatsStatement:
		return s.executeShowStatistics(stmt)
	default:
		panic(fmt.Sprintf("unsupported statement type: %T", stmt))
	}
}

// executeShowStatistics returns the statistics of the registered monitor client in
// the standard form expected by users of the InfluxDB system.
func (s *Service) executeShowStatistics(q *influxql.ShowStatsStatement) *influxql.Result {
	stats, _ := s.statistics()
	rows := make([]*influxql.Row, len(stats))

	for n, stat := range stats {
		row := &influxql.Row{Name: stat.Name}
		values := make([]interface{}, 0, len(stat.Tags)+len(stat.Values))

		for _, k := range stat.tagNames() {
			row.Columns = append(row.Columns, k)
			values = append(values, stat.Tags[k])
		}
		for _, k := range stat.valueNames() {
			row.Columns = append(row.Columns, k)
			values = append(values, stat.Values[k])
		}
		row.Values = [][]interface{}{values}
		rows[n] = row
	}
	return &influxql.Result{Series: rows}
}

// statistics returns the combined statistics for all registered clients.
func (s *Service) statistics() ([]*statistic, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	statistics := make([]*statistic, 0, len(s.registrations))
	for _, r := range s.registrations {
		stats, err := r.Client.Statistics()
		if err != nil {
			continue
		}

		// If a registered client has no field data, don't include it in the results
		if len(stats) == 0 {
			continue
		}

		statistics = append(statistics, newStatistic(r.name, r.tags, stats))
	}
	return statistics, nil
}

// storeStatistics writes the statistics to an InfluxDB system.
func (s *Service) storeStatistics() {
	// XXX add tags such as local hostname and cluster ID
	//a.Tags["clusterID"] = strconv.FormatUint(s.clusterID, 10)
	//a.Tags["nodeID"] = strconv.FormatUint(s.nodeID, 10)
	//a.Tags["hostname"] = s.hostname
	defer s.wg.Done()

	tick := time.NewTicker(s.storeInterval)
	defer tick.Stop()
	for {
		select {
		case <-tick.C:
			// Write stats here.
		case <-s.done:
			s.Logger.Printf("terminating storage of statistics to %s", s.storeAddress)
			return
		}

	}
}

// statistic represents the information returned by a single monitor client.
type statistic struct {
	Name   string
	Tags   map[string]string
	Values map[string]interface{}
}

// newStatistic returns a new statistic object. It ensures that tags are always non-nil.
func newStatistic(name string, tags map[string]string, values map[string]interface{}) *statistic {
	a := tags
	if a == nil {
		a = make(map[string]string)
	}

	return &statistic{
		Name:   name,
		Tags:   a,
		Values: values,
	}
}

// tagNames returns a sorted list of the tag names, if any.
func (s *statistic) tagNames() []string {
	a := make([]string, 0, len(s.Tags))
	for k, _ := range s.Tags {
		a = append(a, k)
	}
	sort.Strings(a)
	return a
}

// valueNames returns a sorted list of the value names, if any.
func (s *statistic) valueNames() []string {
	a := make([]string, 0, len(s.Values))
	for k, _ := range s.Values {
		a = append(a, k)
	}
	sort.Strings(a)
	return a
}

// clientWithMeta wraps a registered client with its associated name and tags.
type clientWithMeta struct {
	Client
	name string
	tags map[string]string
}

// MonitorClient wraps a *expvar.Map so that it implements the Client interface. It is for
// use by external packages that just record stats in an expvar.Map type.
type MonitorClient struct {
	ep *expvar.Map
}

// NewMonitorClient returns a new MonitorClient using the given expvar.Map.
func NewMonitorClient(ep *expvar.Map) *MonitorClient {
	return &MonitorClient{ep: ep}
}

// Statistics implements the Client interface for a MonitorClient.
func (m MonitorClient) Statistics() (map[string]interface{}, error) {
	values := make(map[string]interface{})
	m.ep.Do(func(kv expvar.KeyValue) {
		var f interface{}
		var err error
		switch v := kv.Value.(type) {
		case *expvar.Float:
			f, err = strconv.ParseFloat(v.String(), 64)
			if err != nil {
				return
			}
		case *expvar.Int:
			f, err = strconv.ParseUint(v.String(), 10, 64)
			if err != nil {
				return
			}
		default:
			return
		}
		values[kv.Key] = f
	})

	return values, nil
}

// Diagnostics implements the Client interface for a MonitorClient.
func (m MonitorClient) Diagnostics() (map[string]interface{}, error) {
	return nil, nil
}

func ensureDatabaseExists(host, database string) error {
	values := url.Values{}
	values.Set("q", fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", database))
	resp, err := http.Get(host + "/query?" + values.Encode())
	if err != nil {
		return fmt.Errorf("failed to create monitoring database on %s: %s", host, err.Error())
	}
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("failed to create monitoring database on %s, received code: %d",
			host, resp.StatusCode)
	}
	return nil
}
<commit_msg>Set tags correctly for SHOW STATS<commit_after>package monitor

import (
	"expvar"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"sort"
	"strconv"
	"sync"
	"time"

	"github.com/influxdb/influxdb/influxql"
)

// Client is the interface modules must implement if they wish to register with monitor.
type Client interface {
	Statistics() (map[string]interface{}, error)
	Diagnostics() (map[string]interface{}, error)
}

// Service represents an instance of the monitor service.
type Service struct {
	wg            sync.WaitGroup
	done          chan struct{}
	mu            sync.Mutex
	registrations []*clientWithMeta

	hostname  string
	clusterID uint64
	nodeID    uint64

	storeEnabled  bool
	storeDatabase string
	storeAddress  string
	storeInterval time.Duration

	Logger *log.Logger
}

// NewService returns a new instance of the monitor service.
func NewService(c Config) *Service {
	return &Service{
		registrations: make([]*clientWithMeta, 0),
		storeEnabled:  c.StoreEnabled,
		storeDatabase: c.StoreDatabase,
		storeAddress:  c.StoreAddress,
		storeInterval: time.Duration(c.StoreInterval),
		Logger:        log.New(os.Stderr, "[monitor] ", log.LstdFlags),
	}
}

// Open opens the monitoring service, using the given clusterID, node ID, and hostname
// for identification purposes.
func (s *Service) Open(clusterID, nodeID uint64, hostname string) error {
	s.Logger.Printf("starting monitor service for cluster %d, host %s", clusterID, hostname)
	s.clusterID = clusterID
	s.nodeID = nodeID
	s.hostname = hostname

	// Self-register Go runtime stats.
	s.Register("runtime", nil, &goRuntime{})

	// If enabled, record stats in a InfluxDB system.
	if s.storeEnabled {
		s.Logger.Printf("storing in %s, database '%s', interval %s",
			s.storeAddress, s.storeDatabase, s.storeInterval)

		s.Logger.Printf("ensuring database %s exists on %s", s.storeDatabase, s.storeAddress)
		if err := ensureDatabaseExists(s.storeAddress, s.storeDatabase); err != nil {
			return err
		}

		// Start periodic writes to system.
		s.wg.Add(1)
		go s.storeStatistics()
	}

	return nil
}

// Close closes the monitor service.
func (s *Service) Close() {
	s.Logger.Println("shutting down monitor service")
	close(s.done)
	s.wg.Wait()
	s.done = nil
}

// SetLogger sets the internal logger to the logger passed in.
func (s *Service) SetLogger(l *log.Logger) {
	s.Logger = l
}

// Register registers a client with the given name and tags.
func (s *Service) Register(name string, tags map[string]string, client Client) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	c := &clientWithMeta{
		Client: client,
		name:   name,
		tags:   tags,
	}
	s.registrations = append(s.registrations, c)
	s.Logger.Printf(`'%s:%v' registered for monitoring`, name, tags)
	return nil
}

// ExecuteStatement executes monitor-related query statements.
func (s *Service) ExecuteStatement(stmt influxql.Statement) *influxql.Result {
	switch stmt := stmt.(type) {
	case *influxql.ShowStatsStatement:
		return s.executeShowStatistics(stmt)
	default:
		panic(fmt.Sprintf("unsupported statement type: %T", stmt))
	}
}

// executeShowStatistics returns the statistics of the registered monitor client in
// the standard form expected by users of the InfluxDB system.
func (s *Service) executeShowStatistics(q *influxql.ShowStatsStatement) *influxql.Result {
	stats, _ := s.statistics()
	rows := make([]*influxql.Row, len(stats))

	for n, stat := range stats {
		row := &influxql.Row{Name: stat.Name, Tags: stat.Tags}

		values := make([]interface{}, 0, len(stat.Values))
		for _, k := range stat.valueNames() {
			row.Columns = append(row.Columns, k)
			values = append(values, stat.Values[k])
		}
		row.Values = [][]interface{}{values}
		rows[n] = row
	}
	return &influxql.Result{Series: rows}
}

// statistics returns the combined statistics for all registered clients.
func (s *Service) statistics() ([]*statistic, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	statistics := make([]*statistic, 0, len(s.registrations))
	for _, r := range s.registrations {
		stats, err := r.Client.Statistics()
		if err != nil {
			continue
		}

		// If a registered client has no field data, don't include it in the results
		if len(stats) == 0 {
			continue
		}

		statistics = append(statistics, newStatistic(r.name, r.tags, stats))
	}
	return statistics, nil
}

// storeStatistics writes the statistics to an InfluxDB system.
func (s *Service) storeStatistics() {
	// XXX add tags such as local hostname and cluster ID
	//a.Tags["clusterID"] = strconv.FormatUint(s.clusterID, 10)
	//a.Tags["nodeID"] = strconv.FormatUint(s.nodeID, 10)
	//a.Tags["hostname"] = s.hostname
	defer s.wg.Done()

	tick := time.NewTicker(s.storeInterval)
	defer tick.Stop()
	for {
		select {
		case <-tick.C:
			// Write stats here.
		case <-s.done:
			s.Logger.Printf("terminating storage of statistics to %s", s.storeAddress)
			return
		}

	}
}

// statistic represents the information returned by a single monitor client.
type statistic struct {
	Name   string
	Tags   map[string]string
	Values map[string]interface{}
}

// newStatistic returns a new statistic object. It ensures that tags are always non-nil.
func newStatistic(name string, tags map[string]string, values map[string]interface{}) *statistic {
	a := tags
	if a == nil {
		a = make(map[string]string)
	}

	return &statistic{
		Name:   name,
		Tags:   a,
		Values: values,
	}
}

// tagNames returns a sorted list of the tag names, if any.
func (s *statistic) tagNames() []string {
	a := make([]string, 0, len(s.Tags))
	for k, _ := range s.Tags {
		a = append(a, k)
	}
	sort.Strings(a)
	return a
}

// valueNames returns a sorted list of the value names, if any.
func (s *statistic) valueNames() []string {
	a := make([]string, 0, len(s.Values))
	for k, _ := range s.Values {
		a = append(a, k)
	}
	sort.Strings(a)
	return a
}

// clientWithMeta wraps a registered client with its associated name and tags.
type clientWithMeta struct {
	Client
	name string
	tags map[string]string
}

// MonitorClient wraps a *expvar.Map so that it implements the Client interface. It is for
// use by external packages that just record stats in an expvar.Map type.
type MonitorClient struct {
	ep *expvar.Map
}

// NewMonitorClient returns a new MonitorClient using the given expvar.Map.
func NewMonitorClient(ep *expvar.Map) *MonitorClient {
	return &MonitorClient{ep: ep}
}

// Statistics implements the Client interface for a MonitorClient.
func (m MonitorClient) Statistics() (map[string]interface{}, error) {
	values := make(map[string]interface{})
	m.ep.Do(func(kv expvar.KeyValue) {
		var f interface{}
		var err error
		switch v := kv.Value.(type) {
		case *expvar.Float:
			f, err = strconv.ParseFloat(v.String(), 64)
			if err != nil {
				return
			}
		case *expvar.Int:
			f, err = strconv.ParseUint(v.String(), 10, 64)
			if err != nil {
				return
			}
		default:
			return
		}
		values[kv.Key] = f
	})

	return values, nil
}

// Diagnostics implements the Client interface for a MonitorClient.
func (m MonitorClient) Diagnostics() (map[string]interface{}, error) {
	return nil, nil
}

func ensureDatabaseExists(host, database string) error {
	values := url.Values{}
	values.Set("q", fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", database))
	resp, err := http.Get(host + "/query?" + values.Encode())
	if err != nil {
		return fmt.Errorf("failed to create monitoring database on %s: %s", host, err.Error())
	}
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("failed to create monitoring database on %s, received code: %d",
			host, resp.StatusCode)
	}
	return nil
}
<|endoftext|>


<commit_before>package netlink

import (
	"fmt"
	"net"
	"strings"
	"syscall"

	"github.com/vishvananda/netlink/nl"
	"github.com/vishvananda/netns"
	"golang.org/x/sys/unix"
)

// AddrAdd will add an IP address to a link device.
//
// Equivalent to: `ip addr add $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func AddrAdd(link Link, addr *Addr) error {
	return pkgHandle.AddrAdd(link, addr)
}

// AddrAdd will add an IP address to a link device.
//
// Equivalent to: `ip addr add $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func (h *Handle) AddrAdd(link Link, addr *Addr) error {
	req := h.newNetlinkRequest(unix.RTM_NEWADDR, unix.NLM_F_CREATE|unix.NLM_F_EXCL|unix.NLM_F_ACK)
	return h.addrHandle(link, addr, req)
}

// AddrReplace will replace (or, if not present, add) an IP address on a link device.
//
// Equivalent to: `ip addr replace $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func AddrReplace(link Link, addr *Addr) error {
	return pkgHandle.AddrReplace(link, addr)
}

// AddrReplace will replace (or, if not present, add) an IP address on a link device.
//
// Equivalent to: `ip addr replace $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func (h *Handle) AddrReplace(link Link, addr *Addr) error {
	req := h.newNetlinkRequest(unix.RTM_NEWADDR, unix.NLM_F_CREATE|unix.NLM_F_REPLACE|unix.NLM_F_ACK)
	return h.addrHandle(link, addr, req)
}

// AddrDel will delete an IP address from a link device.
//
// Equivalent to: `ip addr del $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func AddrDel(link Link, addr *Addr) error {
	return pkgHandle.AddrDel(link, addr)
}

// AddrDel will delete an IP address from a link device.
// Equivalent to: `ip addr del $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func (h *Handle) AddrDel(link Link, addr *Addr) error {
	req := h.newNetlinkRequest(unix.RTM_DELADDR, unix.NLM_F_ACK)
	return h.addrHandle(link, addr, req)
}

func (h *Handle) addrHandle(link Link, addr *Addr, req *nl.NetlinkRequest) error {
	base := link.Attrs()
	if addr.Label != "" && !strings.HasPrefix(addr.Label, base.Name) {
		return fmt.Errorf("label must begin with interface name")
	}
	h.ensureIndex(base)

	family := nl.GetIPFamily(addr.IP)

	msg := nl.NewIfAddrmsg(family)
	msg.Index = uint32(base.Index)
	msg.Scope = uint8(addr.Scope)
	mask := addr.Mask
	if addr.Peer != nil {
		mask = addr.Peer.Mask
	}
	prefixlen, masklen := mask.Size()
	msg.Prefixlen = uint8(prefixlen)
	req.AddData(msg)

	var localAddrData []byte
	if family == FAMILY_V4 {
		localAddrData = addr.IP.To4()
	} else {
		localAddrData = addr.IP.To16()
	}

	localData := nl.NewRtAttr(unix.IFA_LOCAL, localAddrData)
	req.AddData(localData)
	var peerAddrData []byte
	if addr.Peer != nil {
		if family == FAMILY_V4 {
			peerAddrData = addr.Peer.IP.To4()
		} else {
			peerAddrData = addr.Peer.IP.To16()
		}
	} else {
		peerAddrData = localAddrData
	}

	addressData := nl.NewRtAttr(unix.IFA_ADDRESS, peerAddrData)
	req.AddData(addressData)

	if addr.Flags != 0 {
		if addr.Flags <= 0xff {
			msg.IfAddrmsg.Flags = uint8(addr.Flags)
		} else {
			b := make([]byte, 4)
			native.PutUint32(b, uint32(addr.Flags))
			flagsData := nl.NewRtAttr(unix.IFA_FLAGS, b)
			req.AddData(flagsData)
		}
	}

	if family == FAMILY_V4 {
		// Automatically set the broadcast address if it is unset and the
		// subnet is large enough to sensibly have one (/30 or larger).
		// See: RFC 3021
		if addr.Broadcast == nil && prefixlen < 31 {
			calcBroadcast := make(net.IP, masklen/8)
			for i := range localAddrData {
				calcBroadcast[i] = localAddrData[i] | ^mask[i]
			}
			addr.Broadcast = calcBroadcast
		}

		if addr.Broadcast != nil {
			req.AddData(nl.NewRtAttr(unix.IFA_BROADCAST, addr.Broadcast))
		}

		if addr.Label != "" {
			labelData := nl.NewRtAttr(unix.IFA_LABEL, nl.ZeroTerminated(addr.Label))
			req.AddData(labelData)
		}
	}

	// 0 is the default value for these attributes. However, 0 means "expired", while the least-surprising default
	// value should be "forever". To compensate for that, only add the attributes if at least one of the values is
	// non-zero, which means the caller has explicitly set them
	if addr.ValidLft > 0 || addr.PreferedLft > 0 {
		cachedata := nl.IfaCacheInfo{unix.IfaCacheinfo{
			Valid:    uint32(addr.ValidLft),
			Prefered: uint32(addr.PreferedLft),
		}}
		req.AddData(nl.NewRtAttr(unix.IFA_CACHEINFO, cachedata.Serialize()))
	}

	_, err := req.Execute(unix.NETLINK_ROUTE, 0)
	return err
}

// AddrList gets a list of IP addresses in the system.
// Equivalent to: `ip addr show`.
// The list can be filtered by link and ip family.
func AddrList(link Link, family int) ([]Addr, error) {
	return pkgHandle.AddrList(link, family)
}

// AddrList gets a list of IP addresses in the system.
// Equivalent to: `ip addr show`.
// The list can be filtered by link and ip family.
func (h *Handle) AddrList(link Link, family int) ([]Addr, error) {
	req := h.newNetlinkRequest(unix.RTM_GETADDR, unix.NLM_F_DUMP)
	msg := nl.NewIfInfomsg(family)
	req.AddData(msg)

	msgs, err := req.Execute(unix.NETLINK_ROUTE, unix.RTM_NEWADDR)
	if err != nil {
		return nil, err
	}

	indexFilter := 0
	if link != nil {
		base := link.Attrs()
		h.ensureIndex(base)
		indexFilter = base.Index
	}

	var res []Addr
	for _, m := range msgs {
		addr, msgFamily, err := parseAddr(m)
		if err != nil {
			return res, err
		}

		if link != nil && addr.LinkIndex != indexFilter {
			// Ignore messages from other interfaces
			continue
		}

		if family != FAMILY_ALL && msgFamily != family {
			continue
		}

		res = append(res, addr)
	}

	return res, nil
}

func parseAddr(m []byte) (addr Addr, family int, err error) {
	msg := nl.DeserializeIfAddrmsg(m)

	family = -1
	addr.LinkIndex = -1

	attrs, err1 := nl.ParseRouteAttr(m[msg.Len():])
	if err1 != nil {
		err = err1
		return
	}

	family = int(msg.Family)
	addr.LinkIndex = int(msg.Index)

	var local, dst *net.IPNet
	for _, attr := range attrs {
		switch attr.Attr.Type {
		case unix.IFA_ADDRESS:
			dst = &net.IPNet{
				IP:   attr.Value,
				Mask: net.CIDRMask(int(msg.Prefixlen), 8*len(attr.Value)),
			}
		case unix.IFA_LOCAL:
			// iproute2 manual:
			// If a peer address is specified, the local address
			// cannot have a prefix length. The network prefix is
			// associated with the peer rather than with the local
			// address.
			n := 8 * len(attr.Value)
			local = &net.IPNet{
				IP:   attr.Value,
				Mask: net.CIDRMask(n, n),
			}
		case unix.IFA_BROADCAST:
			addr.Broadcast = attr.Value
		case unix.IFA_LABEL:
			addr.Label = string(attr.Value[:len(attr.Value)-1])
		case unix.IFA_FLAGS:
			addr.Flags = int(native.Uint32(attr.Value[0:4]))
		case unix.IFA_CACHEINFO:
			ci := nl.DeserializeIfaCacheInfo(attr.Value)
			addr.PreferedLft = int(ci.Prefered)
			addr.ValidLft = int(ci.Valid)
		}
	}

	// libnl addr.c comment:
	// IPv6 sends the local address as IFA_ADDRESS with no
	// IFA_LOCAL, IPv4 sends both IFA_LOCAL and IFA_ADDRESS
	// with IFA_ADDRESS being the peer address if they differ
	//
	// But obviously, as there are IPv6 PtP addresses, too,
	// IFA_LOCAL should also be handled for IPv6.
	if local != nil {
		if family == FAMILY_V4 && dst != nil && local.IP.Equal(dst.IP) {
			addr.IPNet = dst
		} else {
			addr.IPNet = local
			addr.Peer = dst
		}
	} else {
		addr.IPNet = dst
	}

	addr.Scope = int(msg.Scope)

	return
}

type AddrUpdate struct {
	LinkAddress net.IPNet
	LinkIndex   int
	Flags       int
	Scope       int
	PreferedLft int
	ValidLft    int
	NewAddr     bool // true=added false=deleted
}

// AddrSubscribe takes a chan down which notifications will be sent
// when addresses change.  Close the 'done' chan to stop subscription.
func AddrSubscribe(ch chan<- AddrUpdate, done <-chan struct{}) error {
	return addrSubscribeAt(netns.None(), netns.None(), ch, done, nil, false, 0)
}

// AddrSubscribeAt works like AddrSubscribe plus it allows the caller
// to choose the network namespace in which to subscribe (ns).
func AddrSubscribeAt(ns netns.NsHandle, ch chan<- AddrUpdate, done <-chan struct{}) error {
	return addrSubscribeAt(ns, netns.None(), ch, done, nil, false, 0)
}

// AddrSubscribeOptions contains a set of options to use with
// AddrSubscribeWithOptions.
type AddrSubscribeOptions struct {
	Namespace         *netns.NsHandle
	ErrorCallback     func(error)
	ListExisting      bool
	ReceiveBufferSize int
}

// AddrSubscribeWithOptions work like AddrSubscribe but enable to
// provide additional options to modify the behavior. Currently, the
// namespace can be provided as well as an error callback.
func AddrSubscribeWithOptions(ch chan<- AddrUpdate, done <-chan struct{}, options AddrSubscribeOptions) error {
	if options.Namespace == nil {
		none := netns.None()
		options.Namespace = &none
	}
	return addrSubscribeAt(*options.Namespace, netns.None(), ch, done, options.ErrorCallback, options.ListExisting, options.ReceiveBufferSize)
}

func addrSubscribeAt(newNs, curNs netns.NsHandle, ch chan<- AddrUpdate, done <-chan struct{}, cberr func(error), listExisting bool, rcvbuf int) error {
	s, err := nl.SubscribeAt(newNs, curNs, unix.NETLINK_ROUTE, unix.RTNLGRP_IPV4_IFADDR, unix.RTNLGRP_IPV6_IFADDR)
	if err != nil {
		return err
	}
	if done != nil {
		go func() {
			<-done
			s.Close()
		}()
	}
	if rcvbuf != 0 {
		err = pkgHandle.SetSocketReceiveBufferSize(rcvbuf, false)
		if err != nil {
			return err
		}
	}
	if listExisting {
		req := pkgHandle.newNetlinkRequest(unix.RTM_GETADDR,
			unix.NLM_F_DUMP)
		infmsg := nl.NewIfInfomsg(unix.AF_UNSPEC)
		req.AddData(infmsg)
		if err := s.Send(req); err != nil {
			return err
		}
	}
	go func() {
		defer close(ch)
		for {
			msgs, from, err := s.Receive()
			if err != nil {
				if cberr != nil {
					cberr(fmt.Errorf("Receive failed: %v",
						err))
				}
				return
			}
			if from.Pid != nl.PidKernel {
				if cberr != nil {
					cberr(fmt.Errorf("Wrong sender portid %d, expected %d", from.Pid, nl.PidKernel))
				}
				continue
			}
			for _, m := range msgs {
				if m.Header.Type == unix.NLMSG_DONE {
					continue
				}
				if m.Header.Type == unix.NLMSG_ERROR {
					error := int32(native.Uint32(m.Data[0:4]))
					if error == 0 {
						continue
					}
					if cberr != nil {
						cberr(fmt.Errorf("error message: %v",
							syscall.Errno(-error)))
					}
					continue
				}
				msgType := m.Header.Type
				if msgType != unix.RTM_NEWADDR && msgType != unix.RTM_DELADDR {
					if cberr != nil {
						cberr(fmt.Errorf("bad message type: %d", msgType))
					}
					continue
				}

				addr, _, err := parseAddr(m.Data)
				if err != nil {
					if cberr != nil {
						cberr(fmt.Errorf("could not parse address: %v", err))
					}
					continue
				}

				ch <- AddrUpdate{LinkAddress: *addr.IPNet,
					LinkIndex:   addr.LinkIndex,
					NewAddr:     msgType == unix.RTM_NEWADDR,
					Flags:       addr.Flags,
					Scope:       addr.Scope,
					PreferedLft: addr.PreferedLft,
					ValidLft:    addr.ValidLft}
			}
		}
	}()

	return nil
}
<commit_msg>addr: add receive timeout option for AddrSubscribe<commit_after>package netlink

import (
	"fmt"
	"net"
	"strings"
	"syscall"

	"github.com/vishvananda/netlink/nl"
	"github.com/vishvananda/netns"
	"golang.org/x/sys/unix"
)

// AddrAdd will add an IP address to a link device.
//
// Equivalent to: `ip addr add $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func AddrAdd(link Link, addr *Addr) error {
	return pkgHandle.AddrAdd(link, addr)
}

// AddrAdd will add an IP address to a link device.
//
// Equivalent to: `ip addr add $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func (h *Handle) AddrAdd(link Link, addr *Addr) error {
	req := h.newNetlinkRequest(unix.RTM_NEWADDR, unix.NLM_F_CREATE|unix.NLM_F_EXCL|unix.NLM_F_ACK)
	return h.addrHandle(link, addr, req)
}

// AddrReplace will replace (or, if not present, add) an IP address on a link device.
//
// Equivalent to: `ip addr replace $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func AddrReplace(link Link, addr *Addr) error {
	return pkgHandle.AddrReplace(link, addr)
}

// AddrReplace will replace (or, if not present, add) an IP address on a link device.
//
// Equivalent to: `ip addr replace $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func (h *Handle) AddrReplace(link Link, addr *Addr) error {
	req := h.newNetlinkRequest(unix.RTM_NEWADDR, unix.NLM_F_CREATE|unix.NLM_F_REPLACE|unix.NLM_F_ACK)
	return h.addrHandle(link, addr, req)
}

// AddrDel will delete an IP address from a link device.
//
// Equivalent to: `ip addr del $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func AddrDel(link Link, addr *Addr) error {
	return pkgHandle.AddrDel(link, addr)
}

// AddrDel will delete an IP address from a link device.
// Equivalent to: `ip addr del $addr dev $link`
//
// If `addr` is an IPv4 address and the broadcast address is not given, it
// will be automatically computed based on the IP mask if /30 or larger.
func (h *Handle) AddrDel(link Link, addr *Addr) error {
	req := h.newNetlinkRequest(unix.RTM_DELADDR, unix.NLM_F_ACK)
	return h.addrHandle(link, addr, req)
}

func (h *Handle) addrHandle(link Link, addr *Addr, req *nl.NetlinkRequest) error {
	base := link.Attrs()
	if addr.Label != "" && !strings.HasPrefix(addr.Label, base.Name) {
		return fmt.Errorf("label must begin with interface name")
	}
	h.ensureIndex(base)

	family := nl.GetIPFamily(addr.IP)

	msg := nl.NewIfAddrmsg(family)
	msg.Index = uint32(base.Index)
	msg.Scope = uint8(addr.Scope)
	mask := addr.Mask
	if addr.Peer != nil {
		mask = addr.Peer.Mask
	}
	prefixlen, masklen := mask.Size()
	msg.Prefixlen = uint8(prefixlen)
	req.AddData(msg)

	var localAddrData []byte
	if family == FAMILY_V4 {
		localAddrData = addr.IP.To4()
	} else {
		localAddrData = addr.IP.To16()
	}

	localData := nl.NewRtAttr(unix.IFA_LOCAL, localAddrData)
	req.AddData(localData)
	var peerAddrData []byte
	if addr.Peer != nil {
		if family == FAMILY_V4 {
			peerAddrData = addr.Peer.IP.To4()
		} else {
			peerAddrData = addr.Peer.IP.To16()
		}
	} else {
		peerAddrData = localAddrData
	}

	addressData := nl.NewRtAttr(unix.IFA_ADDRESS, peerAddrData)
	req.AddData(addressData)

	if addr.Flags != 0 {
		if addr.Flags <= 0xff {
			msg.IfAddrmsg.Flags = uint8(addr.Flags)
		} else {
			b := make([]byte, 4)
			native.PutUint32(b, uint32(addr.Flags))
			flagsData := nl.NewRtAttr(unix.IFA_FLAGS, b)
			req.AddData(flagsData)
		}
	}

	if family == FAMILY_V4 {
		// Automatically set the broadcast address if it is unset and the
		// subnet is large enough to sensibly have one (/30 or larger).
		// See: RFC 3021
		if addr.Broadcast == nil && prefixlen < 31 {
			calcBroadcast := make(net.IP, masklen/8)
			for i := range localAddrData {
				calcBroadcast[i] = localAddrData[i] | ^mask[i]
			}
			addr.Broadcast = calcBroadcast
		}

		if addr.Broadcast != nil {
			req.AddData(nl.NewRtAttr(unix.IFA_BROADCAST, addr.Broadcast))
		}

		if addr.Label != "" {
			labelData := nl.NewRtAttr(unix.IFA_LABEL, nl.ZeroTerminated(addr.Label))
			req.AddData(labelData)
		}
	}

	// 0 is the default value for these attributes. However, 0 means "expired", while the least-surprising default
	// value should be "forever". To compensate for that, only add the attributes if at least one of the values is
	// non-zero, which means the caller has explicitly set them
	if addr.ValidLft > 0 || addr.PreferedLft > 0 {
		cachedata := nl.IfaCacheInfo{unix.IfaCacheinfo{
			Valid:    uint32(addr.ValidLft),
			Prefered: uint32(addr.PreferedLft),
		}}
		req.AddData(nl.NewRtAttr(unix.IFA_CACHEINFO, cachedata.Serialize()))
	}

	_, err := req.Execute(unix.NETLINK_ROUTE, 0)
	return err
}

// AddrList gets a list of IP addresses in the system.
// Equivalent to: `ip addr show`.
// The list can be filtered by link and ip family.
func AddrList(link Link, family int) ([]Addr, error) {
	return pkgHandle.AddrList(link, family)
}

// AddrList gets a list of IP addresses in the system.
// Equivalent to: `ip addr show`.
// The list can be filtered by link and ip family.
func (h *Handle) AddrList(link Link, family int) ([]Addr, error) {
	req := h.newNetlinkRequest(unix.RTM_GETADDR, unix.NLM_F_DUMP)
	msg := nl.NewIfInfomsg(family)
	req.AddData(msg)

	msgs, err := req.Execute(unix.NETLINK_ROUTE, unix.RTM_NEWADDR)
	if err != nil {
		return nil, err
	}

	indexFilter := 0
	if link != nil {
		base := link.Attrs()
		h.ensureIndex(base)
		indexFilter = base.Index
	}

	var res []Addr
	for _, m := range msgs {
		addr, msgFamily, err := parseAddr(m)
		if err != nil {
			return res, err
		}

		if link != nil && addr.LinkIndex != indexFilter {
			// Ignore messages from other interfaces
			continue
		}

		if family != FAMILY_ALL && msgFamily != family {
			continue
		}

		res = append(res, addr)
	}

	return res, nil
}

func parseAddr(m []byte) (addr Addr, family int, err error) {
	msg := nl.DeserializeIfAddrmsg(m)

	family = -1
	addr.LinkIndex = -1

	attrs, err1 := nl.ParseRouteAttr(m[msg.Len():])
	if err1 != nil {
		err = err1
		return
	}

	family = int(msg.Family)
	addr.LinkIndex = int(msg.Index)

	var local, dst *net.IPNet
	for _, attr := range attrs {
		switch attr.Attr.Type {
		case unix.IFA_ADDRESS:
			dst = &net.IPNet{
				IP:   attr.Value,
				Mask: net.CIDRMask(int(msg.Prefixlen), 8*len(attr.Value)),
			}
		case unix.IFA_LOCAL:
			// iproute2 manual:
			// If a peer address is specified, the local address
			// cannot have a prefix length. The network prefix is
			// associated with the peer rather than with the local
			// address.
			n := 8 * len(attr.Value)
			local = &net.IPNet{
				IP:   attr.Value,
				Mask: net.CIDRMask(n, n),
			}
		case unix.IFA_BROADCAST:
			addr.Broadcast = attr.Value
		case unix.IFA_LABEL:
			addr.Label = string(attr.Value[:len(attr.Value)-1])
		case unix.IFA_FLAGS:
			addr.Flags = int(native.Uint32(attr.Value[0:4]))
		case unix.IFA_CACHEINFO:
			ci := nl.DeserializeIfaCacheInfo(attr.Value)
			addr.PreferedLft = int(ci.Prefered)
			addr.ValidLft = int(ci.Valid)
		}
	}

	// libnl addr.c comment:
	// IPv6 sends the local address as IFA_ADDRESS with no
	// IFA_LOCAL, IPv4 sends both IFA_LOCAL and IFA_ADDRESS
	// with IFA_ADDRESS being the peer address if they differ
	//
	// But obviously, as there are IPv6 PtP addresses, too,
	// IFA_LOCAL should also be handled for IPv6.
	if local != nil {
		if family == FAMILY_V4 && dst != nil && local.IP.Equal(dst.IP) {
			addr.IPNet = dst
		} else {
			addr.IPNet = local
			addr.Peer = dst
		}
	} else {
		addr.IPNet = dst
	}

	addr.Scope = int(msg.Scope)

	return
}

type AddrUpdate struct {
	LinkAddress net.IPNet
	LinkIndex   int
	Flags       int
	Scope       int
	PreferedLft int
	ValidLft    int
	NewAddr     bool // true=added false=deleted
}

// AddrSubscribe takes a chan down which notifications will be sent
// when addresses change.  Close the 'done' chan to stop subscription.
func AddrSubscribe(ch chan<- AddrUpdate, done <-chan struct{}) error {
	return addrSubscribeAt(netns.None(), netns.None(), ch, done, nil, false, 0, nil)
}

// AddrSubscribeAt works like AddrSubscribe plus it allows the caller
// to choose the network namespace in which to subscribe (ns).
func AddrSubscribeAt(ns netns.NsHandle, ch chan<- AddrUpdate, done <-chan struct{}) error {
	return addrSubscribeAt(ns, netns.None(), ch, done, nil, false, 0, nil)
}

// AddrSubscribeOptions contains a set of options to use with
// AddrSubscribeWithOptions.
type AddrSubscribeOptions struct {
	Namespace         *netns.NsHandle
	ErrorCallback     func(error)
	ListExisting      bool
	ReceiveBufferSize int
	ReceiveTimeout    *unix.Timeval
}

// AddrSubscribeWithOptions work like AddrSubscribe but enable to
// provide additional options to modify the behavior. Currently, the
// namespace can be provided as well as an error callback.
func AddrSubscribeWithOptions(ch chan<- AddrUpdate, done <-chan struct{}, options AddrSubscribeOptions) error {
	if options.Namespace == nil {
		none := netns.None()
		options.Namespace = &none
	}
	return addrSubscribeAt(*options.Namespace, netns.None(), ch, done, options.ErrorCallback, options.ListExisting, options.ReceiveBufferSize, options.ReceiveTimeout)
}

func addrSubscribeAt(newNs, curNs netns.NsHandle, ch chan<- AddrUpdate, done <-chan struct{}, cberr func(error), listExisting bool, rcvbuf int, rcvTimeout *unix.Timeval) error {
	s, err := nl.SubscribeAt(newNs, curNs, unix.NETLINK_ROUTE, unix.RTNLGRP_IPV4_IFADDR, unix.RTNLGRP_IPV6_IFADDR)
	if err != nil {
		return err
	}
	if rcvTimeout != nil {
		if err := s.SetReceiveTimeout(rcvTimeout); err != nil {
			return err
		}
	}

	if done != nil {
		go func() {
			<-done
			s.Close()
		}()
	}
	if rcvbuf != 0 {
		err = pkgHandle.SetSocketReceiveBufferSize(rcvbuf, false)
		if err != nil {
			return err
		}
	}
	if listExisting {
		req := pkgHandle.newNetlinkRequest(unix.RTM_GETADDR,
			unix.NLM_F_DUMP)
		infmsg := nl.NewIfInfomsg(unix.AF_UNSPEC)
		req.AddData(infmsg)
		if err := s.Send(req); err != nil {
			return err
		}
	}
	go func() {
		defer close(ch)
		for {
			msgs, from, err := s.Receive()
			if err != nil {
				if cberr != nil {
					cberr(fmt.Errorf("Receive failed: %v",
						err))
				}
				return
			}
			if from.Pid != nl.PidKernel {
				if cberr != nil {
					cberr(fmt.Errorf("Wrong sender portid %d, expected %d", from.Pid, nl.PidKernel))
				}
				continue
			}
			for _, m := range msgs {
				if m.Header.Type == unix.NLMSG_DONE {
					continue
				}
				if m.Header.Type == unix.NLMSG_ERROR {
					error := int32(native.Uint32(m.Data[0:4]))
					if error == 0 {
						continue
					}
					if cberr != nil {
						cberr(fmt.Errorf("error message: %v",
							syscall.Errno(-error)))
					}
					continue
				}
				msgType := m.Header.Type
				if msgType != unix.RTM_NEWADDR && msgType != unix.RTM_DELADDR {
					if cberr != nil {
						cberr(fmt.Errorf("bad message type: %d", msgType))
					}
					continue
				}

				addr, _, err := parseAddr(m.Data)
				if err != nil {
					if cberr != nil {
						cberr(fmt.Errorf("could not parse address: %v", err))
					}
					continue
				}

				ch <- AddrUpdate{LinkAddress: *addr.IPNet,
					LinkIndex:   addr.LinkIndex,
					NewAddr:     msgType == unix.RTM_NEWADDR,
					Flags:       addr.Flags,
					Scope:       addr.Scope,
					PreferedLft: addr.PreferedLft,
					ValidLft:    addr.ValidLft}
			}
		}
	}()

	return nil
}
<|endoftext|>


<commit_before>// Package run is the core logic of get-headers
package run

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"text/tabwriter"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/carlmjohnson/get-headers/prettyprint"
)

// Don't follow redirects
func checkRedirect(req *http.Request, via []*http.Request) error {
	return http.ErrUseLastResponse
}

// base client for all http requests
var client = http.Client{
	CheckRedirect: checkRedirect,
	Transport: &http.Transport{
		DisableCompression: true,
	},
}

func IPDialer() (*net.Addr, func(ctx context.Context, network, addr string) (net.Conn, error)) {
	var ip net.Addr
	return &ip, func(ctx context.Context, network, addr string) (net.Conn, error) {
		conn, err := net.Dial(network, addr)
		ip = conn.RemoteAddr()
		return conn, err
	}
}

// Main takes a list of urls and request parameters, then fetches the URLs and
// outputs the headers to stdout
func Main(cookie, etag string, gzip, ignoreBody bool, urls ...string) error {
	for i, url := range urls {
		// Separate subsequent lookups with newline
		if i > 0 {
			fmt.Println()
		}
		if err := getHeaders(cookie, etag, gzip, ignoreBody, url); err != nil {
			return err
		}
	}
	return nil
}

func getHeaders(cookie, etag string, gzip, ignoreBody bool, url string) error {
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
	defer cancel()
	eg, ctx := errgroup.WithContext(ctx)

	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return err
	}

	if gzip {
		req.Header.Add("Accept-Encoding", "gzip, deflate")
	}

	if etag != "" {
		req.Header.Add("If-None-Match", etag)
	}

	if cookie != "" {
		req.Header.Add("Cookie", cookie)
	}

	newClient := client
	ip, dc := IPDialer()
	newClient.Transport.(*http.Transport).DialContext = dc
	start := time.Now()
	resp, err := newClient.Do(req)
	duration := time.Since(start)

	if err != nil {
		return err
	}

	var n int64

	if !ignoreBody {
		eg.Go(func() error {
			// Copying to /dev/null just to make sure this is real
			n, err = io.Copy(io.Discard, resp.Body)
			duration = time.Since(start)
			if err != nil {
				return err
			}
			return nil
		})
	}

	fmt.Println("GET", url)
	fmt.Println("Via", *ip)
	fmt.Println(resp.Proto, resp.Status)
	fmt.Println()
	fmt.Println(prettyprint.ResponseHeader(resp.Header))

	if err := eg.Wait(); err != nil {
		return err
	}

	if err := resp.Body.Close(); err != nil {
		return err
	}

	tw := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	fmt.Fprintf(tw, "Time\t%s\n", prettyprint.Duration(duration))
	if n != 0 {
		fmt.Fprintf(tw, "Content length\t%s\n", prettyprint.Size(n))
		bps := prettyprint.Size(float64(n) / duration.Seconds())
		fmt.Fprintf(tw, "Speed\t%s/s\n", bps)
	}
	if err := tw.Flush(); err != nil {
		return err
	}

	return nil
}
<commit_msg>Bug: Fix problem of missing net.Conn on error<commit_after>// Package run is the core logic of get-headers
package run

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"text/tabwriter"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/carlmjohnson/get-headers/prettyprint"
)

// Don't follow redirects
func checkRedirect(req *http.Request, via []*http.Request) error {
	return http.ErrUseLastResponse
}

// base client for all http requests
var client = http.Client{
	CheckRedirect: checkRedirect,
	Transport: &http.Transport{
		DisableCompression: true,
	},
}

func IPDialer() (*net.Addr, func(ctx context.Context, network, addr string) (net.Conn, error)) {
	var ip net.Addr
	return &ip, func(ctx context.Context, network, addr string) (net.Conn, error) {
		conn, err := net.Dial(network, addr)
		if conn != nil {
			ip = conn.RemoteAddr()
		}
		return conn, err
	}
}

// Main takes a list of urls and request parameters, then fetches the URLs and
// outputs the headers to stdout
func Main(cookie, etag string, gzip, ignoreBody bool, urls ...string) error {
	for i, url := range urls {
		// Separate subsequent lookups with newline
		if i > 0 {
			fmt.Println()
		}
		if err := getHeaders(cookie, etag, gzip, ignoreBody, url); err != nil {
			return err
		}
	}
	return nil
}

func getHeaders(cookie, etag string, gzip, ignoreBody bool, url string) error {
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
	defer cancel()
	eg, ctx := errgroup.WithContext(ctx)

	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return err
	}

	if gzip {
		req.Header.Add("Accept-Encoding", "gzip, deflate")
	}

	if etag != "" {
		req.Header.Add("If-None-Match", etag)
	}

	if cookie != "" {
		req.Header.Add("Cookie", cookie)
	}

	newClient := client
	ip, dc := IPDialer()
	newClient.Transport.(*http.Transport).DialContext = dc
	start := time.Now()
	resp, err := newClient.Do(req)
	duration := time.Since(start)

	if err != nil {
		return err
	}

	var n int64

	if !ignoreBody {
		eg.Go(func() error {
			// Copying to /dev/null just to make sure this is real
			n, err = io.Copy(io.Discard, resp.Body)
			duration = time.Since(start)
			if err != nil {
				return err
			}
			return nil
		})
	}

	fmt.Println("GET", url)
	if *ip != nil {
		fmt.Println("Via", *ip)
	}
	fmt.Println(resp.Proto, resp.Status)
	fmt.Println()
	fmt.Println(prettyprint.ResponseHeader(resp.Header))

	if err := eg.Wait(); err != nil {
		return err
	}

	if err := resp.Body.Close(); err != nil {
		return err
	}

	tw := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	fmt.Fprintf(tw, "Time\t%s\n", prettyprint.Duration(duration))
	if n != 0 {
		fmt.Fprintf(tw, "Content length\t%s\n", prettyprint.Size(n))
		bps := prettyprint.Size(float64(n) / duration.Seconds())
		fmt.Fprintf(tw, "Speed\t%s/s\n", bps)
	}
	if err := tw.Flush(); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package go_simple_sql

import (
	"database/sql"
	"fmt"
)

type CONN struct {
	DB  sql.DB
	Err error
}

func (c *CONN) InitDB(ip, port, user, pwd, dbname, charset string) {
	url := user + ":" + pwd + "@" + "tcp(" + ip + ":" + port + ")/" + dbname + "?charset=" + charset
	db, err := sql.Open("mysql", url)
	if err != nil {
		fmt.Println("mysql init fail")
	} else {
		c.DB = *db
		c.Err = err
		fmt.Println("mysql init success")
	}
}

func (c *CONN) Query(text string) ([]map[string]string, error) {
	rows, err := c.DB.Query(text)
	result := make([]map[string]string, 0)
	if err != nil {
		return result, err
	}
	columns, _ := rows.Columns()
	scanArgs := make([]interface{}, len(columns))
	values := make([]interface{}, len(columns))
	for i := range values {
		scanArgs[i] = &values[i]
	}

	for rows.Next() {
		err = rows.Scan(scanArgs...)
		record := make(map[string]string)
		for i, col := range values {
			if col != nil {
				record[columns[i]] = string(col.([]byte))
			}
		}
		result = append(result, record)
	}
	return result, err
}

func (c *CONN) Update(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	rows, err := result.RowsAffected()
	return rows, err
}

func (c *CONN) Insert(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	id, err := result.LastInsertId()
	return id, err
}

func (c *CONN) Delete(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	rows, err := result.RowsAffected()
	return rows, err
}
<commit_msg>fix bug<commit_after>package go_simple_sql

import (
	"database/sql"
	"fmt"
)

type CONN struct {
	DB sql.DB
}

func (c *CONN) InitDB(ip, port, user, pwd, dbname, charset string) error {
	url := user + ":" + pwd + "@" + "tcp(" + ip + ":" + port + ")/" + dbname + "?charset=" + charset
	db, err := sql.Open("mysql", url)
	if err != nil {
		fmt.Println("mysql init fail")
	} else {
		c.DB = *db
		fmt.Println("mysql init success")
	}
}

func (c *CONN) Query(text string) ([]map[string]string, error) {
	rows, err := c.DB.Query(text)
	result := make([]map[string]string, 0)
	if err != nil {
		return result, err
	}
	columns, _ := rows.Columns()
	scanArgs := make([]interface{}, len(columns))
	values := make([]interface{}, len(columns))
	for i := range values {
		scanArgs[i] = &values[i]
	}

	for rows.Next() {
		err = rows.Scan(scanArgs...)
		record := make(map[string]string)
		for i, col := range values {
			if col != nil {
				record[columns[i]] = string(col.([]byte))
			}
		}
		result = append(result, record)
	}
	return result, err
}

func (c *CONN) Update(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	rows, err := result.RowsAffected()
	return rows, err
}

func (c *CONN) Insert(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	id, err := result.LastInsertId()
	return id, err
}

func (c *CONN) Delete(text string) (int64, error) {
	tx, err := c.DB.Begin()
	if err != nil {
		return 0, err
	}
	result, err := tx.Exec(text)
	if err != nil {
		return 0, err
	}
	tx.Commit()
	rows, err := result.RowsAffected()
	return rows, err
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package sshstorage

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"sort"
	"strconv"
	"strings"

	coreerrors "launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/utils"
)

// SSHStorage implements storage.Storage.
//
// The storage is created under sudo, and ownership given over to the
// login uid/gid. This is done so that we don't require sudo, and by
// consequence, don't require a pty, so we can interact with a script
// via stdin.
type SSHStorage struct {
	host       string
	remotepath string
	tmpdir     string

	cmd     *exec.Cmd
	stdin   io.WriteCloser
	stdout  io.ReadCloser
	scanner *bufio.Scanner
}

var sshCommand = func(host string, tty bool, command string) *exec.Cmd {
	sshArgs := []string{host}
	if tty {
		sshArgs = append(sshArgs, "-t")
	}
	sshArgs = append(sshArgs, "--", command)
	return exec.Command("ssh", sshArgs...)
}

type flockmode string

const (
	flockShared    flockmode = "-s"
	flockExclusive flockmode = "-x"
)

// NewSSHStorage creates a new SSHStorage, connected to the
// specified host, managing state under the specified remote path.
//
// A temporary directory must be specified, and should be located on the
// same filesystem as the storage directory to ensure atomic writes.
// The temporary directory will be created when NewSSHStorage is invoked
// if it doesn't already exist; it will never be removed. NewSSHStorage
// will attempt to reassign ownership to the login user, and will return
// an error if it cannot do so.
func NewSSHStorage(host, storagedir, tmpdir string) (*SSHStorage, error) {
	if storagedir == "" {
		return nil, errors.New("storagedir must be specified and non-empty")
	}
	if tmpdir == "" {
		return nil, errors.New("tmpdir must be specified and non-empty")
	}

	script := fmt.Sprintf(
		"install -d -g $SUDO_GID -o $SUDO_UID %s %s",
		utils.ShQuote(storagedir),
		utils.ShQuote(tmpdir),
	)

	cmd := sshCommand(host, true, fmt.Sprintf("sudo bash -c %s", utils.ShQuote(script)))
	cmd.Stdin = os.Stdin
	if out, err := cmd.CombinedOutput(); err != nil {
		err = fmt.Errorf("failed to create storage dir: %v (%v)", err, strings.TrimSpace(string(out)))
		return nil, err
	}

	// We could use sftp, but then we'd be at the mercy of
	// sftp's output messages for checking errors. Instead,
	// we execute an interactive bash shell.
	cmd = sshCommand(host, false, "bash")
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		stdin.Close()
		return nil, err
	}
	stor := &SSHStorage{
		host:       host,
		remotepath: storagedir,
		tmpdir:     tmpdir,
		cmd:        cmd,
		stdin:      stdin,
		stdout:     stdout,
		scanner:    bufio.NewScanner(stdout),
	}
	cmd.Start()

	// Verify we have write permissions.
	_, err = stor.runf(flockExclusive, "touch %s", utils.ShQuote(storagedir))
	if err != nil {
		stdin.Close()
		stdout.Close()
		cmd.Wait()
		return nil, err
	}
	return stor, nil
}

// Close cleanly terminates the underlying SSH connection.
func (s *SSHStorage) Close() error {
	s.stdin.Close()
	s.stdout.Close()
	return s.cmd.Wait()
}

func (s *SSHStorage) runf(flockmode flockmode, command string, args ...interface{}) (string, error) {
	command = fmt.Sprintf(command, args...)
	return s.run(flockmode, command, nil)
}

func (s *SSHStorage) run(flockmode flockmode, command string, input []byte) (string, error) {
	const rcPrefix = "JUJU-RC: "
	command = fmt.Sprintf(
		"SHELL=/bin/bash flock %s %s -c %s",
		flockmode,
		s.remotepath,
		utils.ShQuote(command),
	)
	var encoded string
	if input != nil {
		encoded = base64.StdEncoding.EncodeToString(input)
		command = fmt.Sprintf(
			"head -q -c %d | base64 -d | (%s)",
			len(encoded),
			command,
		)
	}
	command = fmt.Sprintf("(%s) 2>&1; echo %s$?", command, rcPrefix)
	if _, err := s.stdin.Write([]byte(command + "\n")); err != nil {
		return "", fmt.Errorf("failed to write command: %v", err)
	}
	if input != nil {
		if _, err := s.stdin.Write([]byte(encoded)); err != nil {
			return "", fmt.Errorf("failed to write input: %v", err)
		}
	}
	var output []string
	for s.scanner.Scan() {
		line := s.scanner.Text()
		if strings.HasPrefix(line, rcPrefix) {
			line := line[len(rcPrefix):]
			rc, err := strconv.Atoi(line)
			if err != nil {
				return "", fmt.Errorf("failed to parse exit code %q: %v", line, err)
			}
			outputJoined := strings.Join(output, "\n")
			if rc == 0 {
				return outputJoined, nil
			}
			return "", SSHStorageError{outputJoined, rc}
		} else {
			output = append(output, line)
		}
	}
	return "", s.scanner.Err()
}

// path returns a remote absolute path for a storage object name.
func (s *SSHStorage) path(name string) (string, error) {
	remotepath := path.Clean(path.Join(s.remotepath, name))
	if !strings.HasPrefix(remotepath, s.remotepath) {
		return "", fmt.Errorf("%q escapes storage directory", name)
	}
	return remotepath, nil
}

// Get implements storage.StorageReader.Get.
func (s *SSHStorage) Get(name string) (io.ReadCloser, error) {
	path, err := s.path(name)
	if err != nil {
		return nil, err
	}
	out, err := s.runf(flockShared, "base64 < %s", utils.ShQuote(path))
	if err != nil {
		err := err.(SSHStorageError)
		if strings.Contains(err.Output, "No such file") {
			return nil, coreerrors.NewNotFoundError(err, "")
		}
		return nil, err
	}
	decoded, err := base64.StdEncoding.DecodeString(out)
	if err != nil {
		return nil, err
	}
	return ioutil.NopCloser(bytes.NewBuffer(decoded)), nil
}

// List implements storage.StorageReader.List.
func (s *SSHStorage) List(prefix string) ([]string, error) {
	remotepath, err := s.path(prefix)
	if err != nil {
		return nil, err
	}
	dir, prefix := path.Split(remotepath)
	quotedDir := utils.ShQuote(dir)
	out, err := s.runf(flockShared, "(test -d %s && find %s -type f) || true", quotedDir, quotedDir)
	if err != nil {
		return nil, err
	}
	if out == "" {
		return nil, nil
	}
	var names []string
	for _, name := range strings.Split(out, "\n") {
		if strings.HasPrefix(name[len(dir):], prefix) {
			names = append(names, name[len(s.remotepath)+1:])
		}
	}
	sort.Strings(names)
	return names, nil
}

// URL implements storage.StorageReader.URL.
func (s *SSHStorage) URL(name string) (string, error) {
	path, err := s.path(name)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("sftp://%s/%s", s.host, path), nil
}

// DefaultConsistencyStrategy implements storage.StorageReader.ConsistencyStrategy.
func (s *SSHStorage) DefaultConsistencyStrategy() utils.AttemptStrategy {
	return utils.AttemptStrategy{}
}

// ShouldRetry is specified in the StorageReader interface.
func (s *SSHStorage) ShouldRetry(err error) bool {
	return false
}

// Put implements storage.StorageWriter.Put
func (s *SSHStorage) Put(name string, r io.Reader, length int64) error {
	path, err := s.path(name)
	if err != nil {
		return err
	}
	buf := make([]byte, length)
	if _, err := r.Read(buf); err != nil {
		return err
	}
	path = utils.ShQuote(path)
	tmpdir := utils.ShQuote(s.tmpdir)

	// Write to a temporary file ($TMPFILE), then mv atomically.
	command := fmt.Sprintf("mkdir -p `dirname %s` && cat > $TMPFILE", path)
	command = fmt.Sprintf(
		"export TMPDIR=%s && TMPFILE=`mktemp` && ((%s && mv $TMPFILE %s) || rm -f $TMPFILE)",
		tmpdir, command, path,
	)

	_, err = s.run(flockExclusive, command+"\n", buf)
	return err
}

// Remove implements storage.StorageWriter.Remove
func (s *SSHStorage) Remove(name string) error {
	path, err := s.path(name)
	if err != nil {
		return err
	}
	path = utils.ShQuote(path)
	_, err = s.runf(flockExclusive, "rm -f %s", path)
	return err
}

// RemoveAll implements storage.StorageWriter.RemoveAll
func (s *SSHStorage) RemoveAll() error {
	_, err := s.runf(flockExclusive, "rm -fr %s/*", utils.ShQuote(s.remotepath))
	return err
}
<commit_msg>environs/sshstorage: stream stdin<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package sshstorage

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"sort"
	"strconv"
	"strings"

	coreerrors "launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/utils"
)

// SSHStorage implements storage.Storage.
//
// The storage is created under sudo, and ownership given over to the
// login uid/gid. This is done so that we don't require sudo, and by
// consequence, don't require a pty, so we can interact with a script
// via stdin.
type SSHStorage struct {
	host       string
	remotepath string
	tmpdir     string

	cmd     *exec.Cmd
	stdin   io.WriteCloser
	stdout  io.ReadCloser
	scanner *bufio.Scanner
}

var sshCommand = func(host string, tty bool, command string) *exec.Cmd {
	sshArgs := []string{host}
	if tty {
		sshArgs = append(sshArgs, "-t")
	}
	sshArgs = append(sshArgs, "--", command)
	return exec.Command("ssh", sshArgs...)
}

type flockmode string

const (
	flockShared    flockmode = "-s"
	flockExclusive flockmode = "-x"
)

// NewSSHStorage creates a new SSHStorage, connected to the
// specified host, managing state under the specified remote path.
//
// A temporary directory must be specified, and should be located on the
// same filesystem as the storage directory to ensure atomic writes.
// The temporary directory will be created when NewSSHStorage is invoked
// if it doesn't already exist; it will never be removed. NewSSHStorage
// will attempt to reassign ownership to the login user, and will return
// an error if it cannot do so.
func NewSSHStorage(host, storagedir, tmpdir string) (*SSHStorage, error) {
	if storagedir == "" {
		return nil, errors.New("storagedir must be specified and non-empty")
	}
	if tmpdir == "" {
		return nil, errors.New("tmpdir must be specified and non-empty")
	}

	script := fmt.Sprintf(
		"install -d -g $SUDO_GID -o $SUDO_UID %s %s",
		utils.ShQuote(storagedir),
		utils.ShQuote(tmpdir),
	)

	cmd := sshCommand(host, true, fmt.Sprintf("sudo bash -c %s", utils.ShQuote(script)))
	cmd.Stdin = os.Stdin
	if out, err := cmd.CombinedOutput(); err != nil {
		err = fmt.Errorf("failed to create storage dir: %v (%v)", err, strings.TrimSpace(string(out)))
		return nil, err
	}

	// We could use sftp, but then we'd be at the mercy of
	// sftp's output messages for checking errors. Instead,
	// we execute an interactive bash shell.
	cmd = sshCommand(host, false, "bash")
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		stdin.Close()
		return nil, err
	}
	stor := &SSHStorage{
		host:       host,
		remotepath: storagedir,
		tmpdir:     tmpdir,
		cmd:        cmd,
		stdin:      stdin,
		stdout:     stdout,
		scanner:    bufio.NewScanner(stdout),
	}
	cmd.Start()

	// Verify we have write permissions.
	_, err = stor.runf(flockExclusive, "touch %s", utils.ShQuote(storagedir))
	if err != nil {
		stdin.Close()
		stdout.Close()
		cmd.Wait()
		return nil, err
	}
	return stor, nil
}

// Close cleanly terminates the underlying SSH connection.
func (s *SSHStorage) Close() error {
	s.stdin.Close()
	s.stdout.Close()
	return s.cmd.Wait()
}

func (s *SSHStorage) runf(flockmode flockmode, command string, args ...interface{}) (string, error) {
	command = fmt.Sprintf(command, args...)
	return s.run(flockmode, command, nil, 0)
}

func (s *SSHStorage) run(flockmode flockmode, command string, input io.Reader, inputlen int64) (string, error) {
	const rcPrefix = "JUJU-RC: "
	command = fmt.Sprintf(
		"SHELL=/bin/bash flock %s %s -c %s",
		flockmode,
		s.remotepath,
		utils.ShQuote(command),
	)
	if input != nil {
		command = fmt.Sprintf("head -q -n 1 | base64 -d | (%s)", command)
	}
	command = fmt.Sprintf("(%s) 2>&1; echo %s$?", command, rcPrefix)
	if _, err := s.stdin.Write([]byte(command + "\n")); err != nil {
		return "", fmt.Errorf("failed to write command: %v", err)
	}
	if input != nil {
		encoder := base64.NewEncoder(base64.StdEncoding, s.stdin)
		if _, err := io.CopyN(encoder, input, inputlen); err != nil {
			return "", fmt.Errorf("failed to write input: %v", err)
		}
		if err := encoder.Close(); err != nil {
			return "", fmt.Errorf("failed to flush input: %v", err)
		}
		if _, err := s.stdin.Write([]byte("\n")); err != nil {
			return "", fmt.Errorf("failed to terminate input: %v", err)
		}
	}
	var output []string
	for s.scanner.Scan() {
		line := s.scanner.Text()
		if strings.HasPrefix(line, rcPrefix) {
			line := line[len(rcPrefix):]
			rc, err := strconv.Atoi(line)
			if err != nil {
				return "", fmt.Errorf("failed to parse exit code %q: %v", line, err)
			}
			outputJoined := strings.Join(output, "\n")
			if rc == 0 {
				return outputJoined, nil
			}
			return "", SSHStorageError{outputJoined, rc}
		} else {
			output = append(output, line)
		}
	}
	return "", s.scanner.Err()
}

// path returns a remote absolute path for a storage object name.
func (s *SSHStorage) path(name string) (string, error) {
	remotepath := path.Clean(path.Join(s.remotepath, name))
	if !strings.HasPrefix(remotepath, s.remotepath) {
		return "", fmt.Errorf("%q escapes storage directory", name)
	}
	return remotepath, nil
}

// Get implements storage.StorageReader.Get.
func (s *SSHStorage) Get(name string) (io.ReadCloser, error) {
	path, err := s.path(name)
	if err != nil {
		return nil, err
	}
	out, err := s.runf(flockShared, "base64 < %s", utils.ShQuote(path))
	if err != nil {
		err := err.(SSHStorageError)
		if strings.Contains(err.Output, "No such file") {
			return nil, coreerrors.NewNotFoundError(err, "")
		}
		return nil, err
	}
	decoded, err := base64.StdEncoding.DecodeString(out)
	if err != nil {
		return nil, err
	}
	return ioutil.NopCloser(bytes.NewBuffer(decoded)), nil
}

// List implements storage.StorageReader.List.
func (s *SSHStorage) List(prefix string) ([]string, error) {
	remotepath, err := s.path(prefix)
	if err != nil {
		return nil, err
	}
	dir, prefix := path.Split(remotepath)
	quotedDir := utils.ShQuote(dir)
	out, err := s.runf(flockShared, "(test -d %s && find %s -type f) || true", quotedDir, quotedDir)
	if err != nil {
		return nil, err
	}
	if out == "" {
		return nil, nil
	}
	var names []string
	for _, name := range strings.Split(out, "\n") {
		if strings.HasPrefix(name[len(dir):], prefix) {
			names = append(names, name[len(s.remotepath)+1:])
		}
	}
	sort.Strings(names)
	return names, nil
}

// URL implements storage.StorageReader.URL.
func (s *SSHStorage) URL(name string) (string, error) {
	path, err := s.path(name)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("sftp://%s/%s", s.host, path), nil
}

// DefaultConsistencyStrategy implements storage.StorageReader.ConsistencyStrategy.
func (s *SSHStorage) DefaultConsistencyStrategy() utils.AttemptStrategy {
	return utils.AttemptStrategy{}
}

// ShouldRetry is specified in the StorageReader interface.
func (s *SSHStorage) ShouldRetry(err error) bool {
	return false
}

// Put implements storage.StorageWriter.Put
func (s *SSHStorage) Put(name string, r io.Reader, length int64) error {
	path, err := s.path(name)
	if err != nil {
		return err
	}
	path = utils.ShQuote(path)
	tmpdir := utils.ShQuote(s.tmpdir)

	// Write to a temporary file ($TMPFILE), then mv atomically.
	command := fmt.Sprintf("mkdir -p `dirname %s` && cat > $TMPFILE", path)
	command = fmt.Sprintf(
		"TMPFILE=`mktemp --tmpdir=%s` && ((%s && mv $TMPFILE %s) || rm -f $TMPFILE)",
		tmpdir, command, path,
	)

	_, err = s.run(flockExclusive, command+"\n", r, length)
	return err
}

// Remove implements storage.StorageWriter.Remove
func (s *SSHStorage) Remove(name string) error {
	path, err := s.path(name)
	if err != nil {
		return err
	}
	path = utils.ShQuote(path)
	_, err = s.runf(flockExclusive, "rm -f %s", path)
	return err
}

// RemoveAll implements storage.StorageWriter.RemoveAll
func (s *SSHStorage) RemoveAll() error {
	_, err := s.runf(flockExclusive, "rm -fr %s/*", utils.ShQuote(s.remotepath))
	return err
}
<|endoftext|>