<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/libopenstorage/openstorage/api"
	"github.com/libopenstorage/openstorage/config"
	"github.com/libopenstorage/openstorage/volume"
	"github.com/libopenstorage/openstorage/volume/drivers"
)

const (
	// VolumeDriver is the string returned in the handshake protocol.
	VolumeDriver = "VolumeDriver"
)

// Implementation of the Docker volumes plugin specification.
type driver struct {
	restBase
}

type handshakeResp struct {
	Implements []string
}

type volumeRequest struct {
	Name string
	Opts map[string]string
}

type mountRequest struct {
	Name string
	ID   string
}

type volumeResponse struct {
	Err string
}

type volumePathResponse struct {
	Mountpoint string
	volumeResponse
}

type volumeInfo struct {
	Name       string
	Mountpoint string
}

func newVolumePlugin(name string) restServer {
	return &driver{restBase{name: name, version: "0.3"}}
}

func (d *driver) String() string {
	return d.name
}

func volDriverPath(method string) string {
	return fmt.Sprintf("/%s.%s", VolumeDriver, method)
}

func (d *driver) volNotFound(request string, id string, e error, w http.ResponseWriter) error {
	err := fmt.Errorf("Failed to locate volume: " + e.Error())
	d.logRequest(request, id).Warnln(http.StatusNotFound, " ", err.Error())
	return err
}

func (d *driver) volNotMounted(request string, id string) error {
	err := fmt.Errorf("volume not mounted")
	d.logRequest(request, id).Debugln(http.StatusNotFound, " ", err.Error())
	return err
}

func (d *driver) Routes() []*Route {
	return []*Route{
		&Route{verb: "POST", path: volDriverPath("Create"), fn: d.create},
		&Route{verb: "POST", path: volDriverPath("Remove"), fn: d.remove},
		&Route{verb: "POST", path: volDriverPath("Mount"), fn: d.mount},
		&Route{verb: "POST", path: volDriverPath("Path"), fn: d.path},
		&Route{verb: "POST", path: volDriverPath("List"), fn: d.list},
		&Route{verb: "POST", path: volDriverPath("Get"), fn: d.get},
		&Route{verb: "POST", path: volDriverPath("Unmount"), fn: d.unmount},
		&Route{verb: "POST", path: "/Plugin.Activate", fn: d.handshake},
		&Route{verb: "GET", path: "/status", fn: d.status},
	}
}

func (d *driver) emptyResponse(w http.ResponseWriter) {
	json.NewEncoder(w).Encode(&volumeResponse{})
}

func (d *driver) errorResponse(w http.ResponseWriter, err error) {
	json.NewEncoder(w).Encode(&volumeResponse{Err: err.Error()})
}

func (d *driver) volFromName(name string) (*api.Volume, error) {
	v, err := volumedrivers.Get(d.name)
	if err != nil {
		return nil, fmt.Errorf("Cannot locate volume driver for %s: %s", d.name, err.Error())
	}
	vols, err := v.Inspect([]string{name})
	if err == nil && len(vols) == 1 {
		return vols[0], nil
	}
	vols, err = v.Enumerate(&api.VolumeLocator{Name: name}, nil)
	if err == nil && len(vols) == 1 {
		return vols[0], nil
	}
	return nil, fmt.Errorf("Cannot locate volume %s", name)
}

func (d *driver) decode(method string, w http.ResponseWriter, r *http.Request) (*volumeRequest, error) {
	var request volumeRequest
	err := json.NewDecoder(r.Body).Decode(&request)
	if err != nil {
		e := fmt.Errorf("Unable to decode JSON payload")
		d.sendError(method, "", w, e.Error()+":"+err.Error(), http.StatusBadRequest)
		return nil, e
	}
	d.logRequest(method, request.Name).Debugln("")
	return &request, nil
}

func (d *driver) decodeMount(method string, w http.ResponseWriter, r *http.Request) (*mountRequest, error) {
	var request mountRequest
	if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
		e := fmt.Errorf("Unable to decode JSON payload")
		d.sendError(method, "", w, e.Error()+":"+err.Error(), http.StatusBadRequest)
		return nil, e
	}
	d.logRequest(method, request.Name).Debugf("ID: %v", request.ID)
	return &request, nil
}

func (d *driver) handshake(w http.ResponseWriter, r *http.Request) {
	err := json.NewEncoder(w).Encode(&handshakeResp{
		[]string{VolumeDriver},
	})
	if err != nil {
		d.sendError("handshake", "", w, "encode error", http.StatusInternalServerError)
		return
	}
	d.logRequest("handshake", "").Debugln("Handshake completed")
}

func (d *driver) status(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, fmt.Sprintln("osd plugin", d.version))
}

func (d *driver) specFromOpts(Opts map[string]string) *api.VolumeSpec {
	spec := api.VolumeSpec{
		VolumeLabels: make(map[string]string),
		Format:       api.FSType_FS_TYPE_EXT4,
		HaLevel:      1,
	}

	for k, v := range Opts {
		switch k {
		case api.SpecEphemeral:
			spec.Ephemeral, _ = strconv.ParseBool(v)
		case api.SpecSize:
			sizeMulti := uint64(1024 * 1024 * 1024)
			if strings.HasSuffix(v, "G") || strings.HasSuffix(v, "g") {
				sizeMulti = 1024 * 1024 * 1024
				last := len(v) - 1
				v = v[:last]
			}

			size, _ := strconv.ParseUint(v, 10, 64)
			spec.Size = size * sizeMulti
		case api.SpecFilesystem:
			value, _ := api.FSTypeSimpleValueOf(v)
			spec.Format = value
		case api.SpecBlockSize:
			blockSize, _ := strconv.ParseInt(v, 10, 64)
			spec.BlockSize = blockSize
		case api.SpecHaLevel:
			haLevel, _ := strconv.ParseInt(v, 10, 64)
			spec.HaLevel = haLevel
		case api.SpecCos:
			value, _ := strconv.ParseUint(v, 10, 32)
			spec.Cos = uint32(value)
		case api.SpecDedupe:
			spec.Dedupe, _ = strconv.ParseBool(v)
		case api.SpecSnapshotInterval:
			snapshotInterval, _ := strconv.ParseUint(v, 10, 32)
			spec.SnapshotInterval = uint32(snapshotInterval)
		case api.SpecShared:
			shared, _ := strconv.ParseUint(v, 10, 32)
			if shared != 0 {
				spec.Shared = true
			}
		default:
			spec.VolumeLabels[k] = v
		}
	}
	return &spec
}

func (d *driver) mountpath(request *mountRequest) string {
	if len(request.ID) != 0 {
		return path.Join(config.MountBase, request.Name+"_"+request.ID)
	}
	return path.Join(config.MountBase, request.Name)
}

func (d *driver) create(w http.ResponseWriter, r *http.Request) {
	method := "create"
	request, err := d.decode(method, w, r)
	if err != nil {
		return
	}
	d.logRequest(method, request.Name).Infoln("")
	if _, err = d.volFromName(request.Name); err != nil {
		v, err := volumedrivers.Get(d.name)
		if err != nil {
			d.errorResponse(w, err)
			return
		}
		spec := d.specFromOpts(request.Opts)
		if _, err := v.Create(&api.VolumeLocator{Name: request.Name}, nil, spec); err != nil {
			d.errorResponse(w, err)
			return
		}
	}
	json.NewEncoder(w).Encode(&volumeResponse{})
}

func (d *driver) remove(w http.ResponseWriter, r *http.Request) {
	method := "remove"
	request, err := d.decode(method, w, r)
	if err != nil {
		return
	}

	v, err := volumedrivers.Get(d.name)
	if err != nil {
		d.logRequest(method, "").Warnf("Cannot locate volume driver")
		d.errorResponse(w, err)
		return
	}
	if err = v.Delete(request.Name); err != nil {
		d.errorResponse(w, err)
		return
	}
	json.NewEncoder(w).Encode(&volumeResponse{})
}

func (d *driver) mount(w http.ResponseWriter, r *http.Request) {
	var response volumePathResponse
	method := "mount"

	v, err := volumedrivers.Get(d.name)
	if err != nil {
		d.logRequest(method, "").Warnf("Cannot locate volume driver")
		d.errorResponse(w, err)
		return
	}

	request, err := d.decodeMount(method, w, r)
	if err != nil {
		d.errorResponse(w, err)
		return
	}

	vol, err := d.volFromName(request.Name)
	if err != nil {
		d.errorResponse(w, err)
		return
	}

	// If this is a block driver, first attach the volume.
	if v.Type() == api.DriverType_DRIVER_TYPE_BLOCK {
		attachPath, err := v.Attach(vol.Id)
		if err != nil {
			if err == volume.ErrVolAttachedOnRemoteNode {
				d.logRequest(method, request.Name).Infof("Volume is attached on a remote node... will attempt to mount it.")
			} else {
				d.logRequest(method, request.Name).Warnf("Cannot attach volume: %v", err.Error())
				d.errorResponse(w, err)
				return
			}
		} else {
			d.logRequest(method, request.Name).Debugf("response %v", attachPath)
		}
	}

	// Now mount it.
	response.Mountpoint = d.mountpath(request)
	os.MkdirAll(response.Mountpoint, 0755)

	err = v.Mount(vol.Id, response.Mountpoint)
	if err != nil {
		d.logRequest(method, request.Name).Warnf("Cannot mount volume %v, %v",
			response.Mountpoint, err)
		d.errorResponse(w, err)
		return
	}

	d.logRequest(method, request.Name).Infof("response %v", response.Mountpoint)
	json.NewEncoder(w).Encode(&response)
}

func (d *driver) path(w http.ResponseWriter, r *http.Request) {
	method := "path"
	var response volumePathResponse

	request, err := d.decode(method, w, r)
	if err != nil {
		return
	}

	vol, err := d.volFromName(request.Name)
	if err != nil {
		e := d.volNotFound(method, request.Name, err, w)
		d.errorResponse(w, e)
		return
	}

	d.logRequest(method, request.Name).Debugf("")

	if len(vol.AttachPath) == 0 || len(vol.AttachPath) == 0 {
		e := d.volNotMounted(method, request.Name)
		d.errorResponse(w, e)
		return
	}
	response.Mountpoint = vol.AttachPath[0]
	response.Mountpoint = path.Join(response.Mountpoint, config.DataDir)
	d.logRequest(method, request.Name).Debugf("response %v", response.Mountpoint)
	json.NewEncoder(w).Encode(&response)
}

func (d *driver) list(w http.ResponseWriter, r *http.Request) {
	method := "list"

	v, err := volumedrivers.Get(d.name)
	if err != nil {
		d.logRequest(method, "").Warnf("Cannot locate volume driver: %v", err.Error())
		d.errorResponse(w, err)
		return
	}

	vols, err := v.Enumerate(nil, nil)
	if err != nil {
		d.errorResponse(w, err)
		return
	}

	volInfo := make([]volumeInfo, len(vols))
	for i, v := range vols {
		volInfo[i].Name = v.Locator.Name
		if len(v.AttachPath) > 0 || len(v.AttachPath) > 0 {
			volInfo[i].Mountpoint = path.Join(v.AttachPath[0], config.DataDir)
		}
	}
	json.NewEncoder(w).Encode(map[string][]volumeInfo{"Volumes": volInfo})
}

func (d *driver) get(w http.ResponseWriter, r *http.Request) {
	method := "get"

	request, err := d.decode(method, w, r)
	if err != nil {
		return
	}
	vol, err := d.volFromName(request.Name)
	if err != nil {
		e := d.volNotFound(method, request.Name, err, w)
		d.errorResponse(w, e)
		return
	}

	volInfo := volumeInfo{Name: request.Name}
	if len(vol.AttachPath) > 0 || len(vol.AttachPath) > 0 {
		volInfo.Mountpoint = path.Join(vol.AttachPath[0], config.DataDir)
	}

	json.NewEncoder(w).Encode(map[string]volumeInfo{"Volume": volInfo})
}

func (d *driver) unmount(w http.ResponseWriter, r *http.Request) {
	method := "unmount"

	v, err := volumedrivers.Get(d.name)
	if err != nil {
		d.logRequest(method, "").Warnf("Cannot locate volume driver: %v", err.Error())
		d.errorResponse(w, err)
		return
	}

	request, err := d.decodeMount(method, w, r)
	if err != nil {
		return
	}

	vol, err := d.volFromName(request.Name)
	if err != nil {
		e := d.volNotFound(method, request.Name, err, w)
		d.errorResponse(w, e)
		return
	}

	mountpoint := d.mountpath(request)
	err = v.Unmount(vol.Id, mountpoint)
	if err != nil {
		d.logRequest(method, request.Name).Warnf("Cannot unmount volume %v, %v",
			mountpoint, err)
		d.errorResponse(w, err)
		return
	}

	if v.Type() == api.DriverType_DRIVER_TYPE_BLOCK {
		_ = v.Detach(vol.Id)
	}
	d.emptyResponse(w)
}
<commit_msg>Support for docker volume capabilities query<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/libopenstorage/openstorage/api"
	"github.com/libopenstorage/openstorage/config"
	"github.com/libopenstorage/openstorage/volume"
	"github.com/libopenstorage/openstorage/volume/drivers"
)

const (
	// VolumeDriver is the string returned in the handshake protocol.
	VolumeDriver = "VolumeDriver"
)

// Implementation of the Docker volumes plugin specification.
type driver struct {
	restBase
}

type handshakeResp struct {
	Implements []string
}

type volumeRequest struct {
	Name string
	Opts map[string]string
}

type mountRequest struct {
	Name string
	ID   string
}

type volumeResponse struct {
	Err string
}

type volumePathResponse struct {
	Mountpoint string
	volumeResponse
}

type volumeInfo struct {
	Name       string
	Mountpoint string
}

type capabilities struct {
	Scope string
}

type capabilitiesResponse struct {
	Capabilities capabilities
}

func newVolumePlugin(name string) restServer {
	return &driver{restBase{name: name, version: "0.3"}}
}

func (d *driver) String() string {
	return d.name
}

func volDriverPath(method string) string {
	return fmt.Sprintf("/%s.%s", VolumeDriver, method)
}

func (d *driver) volNotFound(request string, id string, e error, w http.ResponseWriter) error {
	err := fmt.Errorf("Failed to locate volume: " + e.Error())
	d.logRequest(request, id).Warnln(http.StatusNotFound, " ", err.Error())
	return err
}

func (d *driver) volNotMounted(request string, id string) error {
	err := fmt.Errorf("volume not mounted")
	d.logRequest(request, id).Debugln(http.StatusNotFound, " ", err.Error())
	return err
}

func (d *driver) Routes() []*Route {
	return []*Route{
		&Route{verb: "POST", path: volDriverPath("Create"), fn: d.create},
		&Route{verb: "POST", path: volDriverPath("Remove"), fn: d.remove},
		&Route{verb: "POST", path: volDriverPath("Mount"), fn: d.mount},
		&Route{verb: "POST", path: volDriverPath("Path"), fn: d.path},
		&Route{verb: "POST", path: volDriverPath("List"), fn: d.list},
		&Route{verb: "POST", path: volDriverPath("Get"), fn: d.get},
		&Route{verb: "POST", path: volDriverPath("Unmount"), fn: d.unmount},
		&Route{verb: "POST", path: volDriverPath("Capabilities"), fn: d.capabilities},
		&Route{verb: "POST", path: "/Plugin.Activate", fn: d.handshake},
		&Route{verb: "GET", path: "/status", fn: d.status},
	}
}

func (d *driver) emptyResponse(w http.ResponseWriter) {
	json.NewEncoder(w).Encode(&volumeResponse{})
}

func (d *driver) errorResponse(w http.ResponseWriter, err error) {
	json.NewEncoder(w).Encode(&volumeResponse{Err: err.Error()})
}

func (d *driver) volFromName(name string) (*api.Volume, error) {
	v, err := volumedrivers.Get(d.name)
	if err != nil {
		return nil, fmt.Errorf("Cannot locate volume driver for %s: %s", d.name, err.Error())
	}
	vols, err := v.Inspect([]string{name})
	if err == nil && len(vols) == 1 {
		return vols[0], nil
	}
	vols, err = v.Enumerate(&api.VolumeLocator{Name: name}, nil)
	if err == nil && len(vols) == 1 {
		return vols[0], nil
	}
	return nil, fmt.Errorf("Cannot locate volume %s", name)
}

func (d *driver) decode(method string, w http.ResponseWriter, r *http.Request) (*volumeRequest, error) {
	var request volumeRequest
	err := json.NewDecoder(r.Body).Decode(&request)
	if err != nil {
		e := fmt.Errorf("Unable to decode JSON payload")
		d.sendError(method, "", w, e.Error()+":"+err.Error(), http.StatusBadRequest)
		return nil, e
	}
	d.logRequest(method, request.Name).Debugln("")
	return &request, nil
}

func (d *driver) decodeMount(method string, w http.ResponseWriter, r *http.Request) (*mountRequest, error) {
	var request mountRequest
	if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
		e := fmt.Errorf("Unable to decode JSON payload")
		d.sendError(method, "", w, e.Error()+":"+err.Error(), http.StatusBadRequest)
		return nil, e
	}
	d.logRequest(method, request.Name).Debugf("ID: %v", request.ID)
	return &request, nil
}

func (d *driver) handshake(w http.ResponseWriter, r *http.Request) {
	err := json.NewEncoder(w).Encode(&handshakeResp{
		[]string{VolumeDriver},
	})
	if err != nil {
		d.sendError("handshake", "", w, "encode error", http.StatusInternalServerError)
		return
	}
	d.logRequest("handshake", "").Debugln("Handshake completed")
}

func (d *driver) status(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, fmt.Sprintln("osd plugin", d.version))
}

func (d *driver) specFromOpts(Opts map[string]string) *api.VolumeSpec {
	spec := api.VolumeSpec{
		VolumeLabels: make(map[string]string),
		Format:       api.FSType_FS_TYPE_EXT4,
		HaLevel:      1,
	}

	for k, v := range Opts {
		switch k {
		case api.SpecEphemeral:
			spec.Ephemeral, _ = strconv.ParseBool(v)
		case api.SpecSize:
			sizeMulti := uint64(1024 * 1024 * 1024)
			if strings.HasSuffix(v, "G") || strings.HasSuffix(v, "g") {
				sizeMulti = 1024 * 1024 * 1024
				last := len(v) - 1
				v = v[:last]
			}

			size, _ := strconv.ParseUint(v, 10, 64)
			spec.Size = size * sizeMulti
		case api.SpecFilesystem:
			value, _ := api.FSTypeSimpleValueOf(v)
			spec.Format = value
		case api.SpecBlockSize:
			blockSize, _ := strconv.ParseInt(v, 10, 64)
			spec.BlockSize = blockSize
		case api.SpecHaLevel:
			haLevel, _ := strconv.ParseInt(v, 10, 64)
			spec.HaLevel = haLevel
		case api.SpecCos:
			value, _ := strconv.ParseUint(v, 10, 32)
			spec.Cos = uint32(value)
		case api.SpecDedupe:
			spec.Dedupe, _ = strconv.ParseBool(v)
		case api.SpecSnapshotInterval:
			snapshotInterval, _ := strconv.ParseUint(v, 10, 32)
			spec.SnapshotInterval = uint32(snapshotInterval)
		case api.SpecShared:
			shared, _ := strconv.ParseUint(v, 10, 32)
			if shared != 0 {
				spec.Shared = true
			}
		default:
			spec.VolumeLabels[k] = v
		}
	}
	return &spec
}

func (d *driver) mountpath(request *mountRequest) string {
	if len(request.ID) != 0 {
		return path.Join(config.MountBase, request.Name+"_"+request.ID)
	}
	return path.Join(config.MountBase, request.Name)
}

func (d *driver) create(w http.ResponseWriter, r *http.Request) {
	method := "create"
	request, err := d.decode(method, w, r)
	if err != nil {
		return
	}
	d.logRequest(method, request.Name).Infoln("")
	if _, err = d.volFromName(request.Name); err != nil {
		v, err := volumedrivers.Get(d.name)
		if err != nil {
			d.errorResponse(w, err)
			return
		}
		spec := d.specFromOpts(request.Opts)
		if _, err := v.Create(&api.VolumeLocator{Name: request.Name}, nil, spec); err != nil {
			d.errorResponse(w, err)
			return
		}
	}
	json.NewEncoder(w).Encode(&volumeResponse{})
}

func (d *driver) remove(w http.ResponseWriter, r *http.Request) {
	method := "remove"
	request, err := d.decode(method, w, r)
	if err != nil {
		return
	}

	v, err := volumedrivers.Get(d.name)
	if err != nil {
		d.logRequest(method, "").Warnf("Cannot locate volume driver")
		d.errorResponse(w, err)
		return
	}
	if err = v.Delete(request.Name); err != nil {
		d.errorResponse(w, err)
		return
	}
	json.NewEncoder(w).Encode(&volumeResponse{})
}

func (d *driver) mount(w http.ResponseWriter, r *http.Request) {
	var response volumePathResponse
	method := "mount"

	v, err := volumedrivers.Get(d.name)
	if err != nil {
		d.logRequest(method, "").Warnf("Cannot locate volume driver")
		d.errorResponse(w, err)
		return
	}

	request, err := d.decodeMount(method, w, r)
	if err != nil {
		d.errorResponse(w, err)
		return
	}

	vol, err := d.volFromName(request.Name)
	if err != nil {
		d.errorResponse(w, err)
		return
	}

	// If this is a block driver, first attach the volume.
	if v.Type() == api.DriverType_DRIVER_TYPE_BLOCK {
		attachPath, err := v.Attach(vol.Id)
		if err != nil {
			if err == volume.ErrVolAttachedOnRemoteNode {
				d.logRequest(method, request.Name).Infof("Volume is attached on a remote node... will attempt to mount it.")
			} else {
				d.logRequest(method, request.Name).Warnf("Cannot attach volume: %v", err.Error())
				d.errorResponse(w, err)
				return
			}
		} else {
			d.logRequest(method, request.Name).Debugf("response %v", attachPath)
		}
	}

	// Now mount it.
	response.Mountpoint = d.mountpath(request)
	os.MkdirAll(response.Mountpoint, 0755)

	err = v.Mount(vol.Id, response.Mountpoint)
	if err != nil {
		d.logRequest(method, request.Name).Warnf("Cannot mount volume %v, %v",
			response.Mountpoint, err)
		d.errorResponse(w, err)
		return
	}

	d.logRequest(method, request.Name).Infof("response %v", response.Mountpoint)
	json.NewEncoder(w).Encode(&response)
}

func (d *driver) path(w http.ResponseWriter, r *http.Request) {
	method := "path"
	var response volumePathResponse

	request, err := d.decode(method, w, r)
	if err != nil {
		return
	}

	vol, err := d.volFromName(request.Name)
	if err != nil {
		e := d.volNotFound(method, request.Name, err, w)
		d.errorResponse(w, e)
		return
	}

	d.logRequest(method, request.Name).Debugf("")

	if len(vol.AttachPath) == 0 || len(vol.AttachPath) == 0 {
		e := d.volNotMounted(method, request.Name)
		d.errorResponse(w, e)
		return
	}
	response.Mountpoint = vol.AttachPath[0]
	response.Mountpoint = path.Join(response.Mountpoint, config.DataDir)
	d.logRequest(method, request.Name).Debugf("response %v", response.Mountpoint)
	json.NewEncoder(w).Encode(&response)
}

func (d *driver) list(w http.ResponseWriter, r *http.Request) {
	method := "list"

	v, err := volumedrivers.Get(d.name)
	if err != nil {
		d.logRequest(method, "").Warnf("Cannot locate volume driver: %v", err.Error())
		d.errorResponse(w, err)
		return
	}

	vols, err := v.Enumerate(nil, nil)
	if err != nil {
		d.errorResponse(w, err)
		return
	}

	volInfo := make([]volumeInfo, len(vols))
	for i, v := range vols {
		volInfo[i].Name = v.Locator.Name
		if len(v.AttachPath) > 0 || len(v.AttachPath) > 0 {
			volInfo[i].Mountpoint = path.Join(v.AttachPath[0], config.DataDir)
		}
	}
	json.NewEncoder(w).Encode(map[string][]volumeInfo{"Volumes": volInfo})
}

func (d *driver) get(w http.ResponseWriter, r *http.Request) {
	method := "get"

	request, err := d.decode(method, w, r)
	if err != nil {
		return
	}
	vol, err := d.volFromName(request.Name)
	if err != nil {
		e := d.volNotFound(method, request.Name, err, w)
		d.errorResponse(w, e)
		return
	}

	volInfo := volumeInfo{Name: request.Name}
	if len(vol.AttachPath) > 0 || len(vol.AttachPath) > 0 {
		volInfo.Mountpoint = path.Join(vol.AttachPath[0], config.DataDir)
	}

	json.NewEncoder(w).Encode(map[string]volumeInfo{"Volume": volInfo})
}

func (d *driver) unmount(w http.ResponseWriter, r *http.Request) {
	method := "unmount"

	v, err := volumedrivers.Get(d.name)
	if err != nil {
		d.logRequest(method, "").Warnf("Cannot locate volume driver: %v", err.Error())
		d.errorResponse(w, err)
		return
	}

	request, err := d.decodeMount(method, w, r)
	if err != nil {
		return
	}

	vol, err := d.volFromName(request.Name)
	if err != nil {
		e := d.volNotFound(method, request.Name, err, w)
		d.errorResponse(w, e)
		return
	}

	mountpoint := d.mountpath(request)
	err = v.Unmount(vol.Id, mountpoint)
	if err != nil {
		d.logRequest(method, request.Name).Warnf("Cannot unmount volume %v, %v",
			mountpoint, err)
		d.errorResponse(w, err)
		return
	}

	if v.Type() == api.DriverType_DRIVER_TYPE_BLOCK {
		_ = v.Detach(vol.Id)
	}
	d.emptyResponse(w)
}

func (d *driver) capabilities(w http.ResponseWriter, r *http.Request) {
	method := "capabilities"
	var response capabilitiesResponse

	response.Capabilities.Scope = "global"
	d.logRequest(method, "").Infof("response %v", response.Capabilities.Scope)
	json.NewEncoder(w).Encode(&response)
}
<|endoftext|>


<commit_before>package nomad

import (
	"crypto/tls"
	"io"
	"net"
	"strings"

	"github.com/armon/go-metrics"
	"github.com/hashicorp/net-rpc-msgpackrpc"
	"github.com/hashicorp/yamux"
)

type RPCType byte

const (
	rpcNomad     RPCType = 0x01
	rpcRaft              = 0x02
	rpcMultiplex         = 0x03
	rpcTLS               = 0x04
)

const (
	// rpcHTTPSMagic is used to detect an incoming HTTPS
	// request. TLS starts with the 0x16 magic byte.
	rpcHTTPSMagic = 0x16

	// rpcHTTPMagic is used to detect an incoming HTTP
	// request. The request starts with 'HTTP'
	rpcHTTPMagic = 0x48
)

// listen is used to listen for incoming RPC connections
func (s *Server) listen() {
	for {
		// Accept a connection
		conn, err := s.rpcListener.Accept()
		if err != nil {
			if s.shutdown {
				return
			}
			s.logger.Printf("[ERR] nomad.rpc: failed to accept RPC conn: %v", err)
			continue
		}

		go s.handleConn(conn, false)
		metrics.IncrCounter([]string{"nomad", "rpc", "accept_conn"}, 1)
	}
}

// handleConn is used to determine if this is a Raft or
// Nomad type RPC connection and invoke the correct handler
func (s *Server) handleConn(conn net.Conn, isTLS bool) {
	// Read a single byte
	buf := make([]byte, 1)
	if _, err := conn.Read(buf); err != nil {
		if err != io.EOF {
			s.logger.Printf("[ERR] nomad.rpc: failed to read byte: %v", err)
		}
		conn.Close()
		return
	}

	// Enforce TLS if VerifyIncoming is set
	if s.config.RequireTLS && !isTLS && RPCType(buf[0]) != rpcTLS {
		s.logger.Printf("[WARN] nomad.rpc: Non-TLS connection attempted with RequireTLS set")
		conn.Close()
		return
	}

	// Switch on the byte
	switch RPCType(buf[0]) {
	case rpcNomad:
		s.handleNomadConn(conn)

	case rpcRaft:
		metrics.IncrCounter([]string{"nomad", "rpc", "raft_handoff"}, 1)
		s.raftLayer.Handoff(conn)

	case rpcMultiplex:
		s.handleMultiplex(conn)

	case rpcTLS:
		if s.rpcTLS == nil {
			s.logger.Printf("[WARN] nomad.rpc: TLS connection attempted, server not configured for TLS")
			conn.Close()
			return
		}
		conn = tls.Server(conn, s.rpcTLS)
		s.handleConn(conn, true)

	default:
		s.logger.Printf("[ERR] nomad.rpc: unrecognized RPC byte: %v", buf[0])
		conn.Close()
		return
	}
}

// handleMultiplex is used to multiplex a single incoming connection
// using the Yamux multiplexer
func (s *Server) handleMultiplex(conn net.Conn) {
	defer conn.Close()
	conf := yamux.DefaultConfig()
	conf.LogOutput = s.config.LogOutput
	server, _ := yamux.Server(conn, conf)
	for {
		sub, err := server.Accept()
		if err != nil {
			if err != io.EOF {
				s.logger.Printf("[ERR] nomad.rpc: multiplex conn accept failed: %v", err)
			}
			return
		}
		go s.handleNomadConn(sub)
	}
}

// handleNomadConn is used to service a single Nomad RPC connection
func (s *Server) handleNomadConn(conn net.Conn) {
	defer conn.Close()
	rpcCodec := msgpackrpc.NewServerCodec(conn)
	for {
		select {
		case <-s.shutdownCh:
			return
		default:
		}

		if err := s.rpcServer.ServeRequest(rpcCodec); err != nil {
			if err != io.EOF && !strings.Contains(err.Error(), "closed") {
				s.logger.Printf("[ERR] nomad.rpc: RPC error: %v (%v)", err, conn)
				metrics.IncrCounter([]string{"nomad", "rpc", "request_error"}, 1)
			}
			return
		}
		metrics.IncrCounter([]string{"nomad", "rpc", "request"}, 1)
	}
}
<commit_msg>nomad: adding raftApply RPC<commit_after>package nomad

import (
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"github.com/armon/go-metrics"
	"github.com/hashicorp/net-rpc-msgpackrpc"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/yamux"
)

type RPCType byte

const (
	rpcNomad     RPCType = 0x01
	rpcRaft              = 0x02
	rpcMultiplex         = 0x03
	rpcTLS               = 0x04
)

const (
	// rpcHTTPSMagic is used to detect an incoming HTTPS
	// request. TLS starts with the 0x16 magic byte.
	rpcHTTPSMagic = 0x16

	// rpcHTTPMagic is used to detect an incoming HTTP
	// request. The request starts with 'HTTP'
	rpcHTTPMagic = 0x48
)

const (
	// Warn if the Raft command is larger than this.
	// If it's over 1MB something is probably being abusive.
	raftWarnSize = 1024 * 1024

	// enqueueLimit caps how long we will wait to enqueue
	// a new Raft command. Something is probably wrong if this
	// value is ever reached. However, it prevents us from blocking
	// the requesting goroutine forever.
	enqueueLimit = 30 * time.Second
)

// listen is used to listen for incoming RPC connections
func (s *Server) listen() {
	for {
		// Accept a connection
		conn, err := s.rpcListener.Accept()
		if err != nil {
			if s.shutdown {
				return
			}
			s.logger.Printf("[ERR] nomad.rpc: failed to accept RPC conn: %v", err)
			continue
		}

		go s.handleConn(conn, false)
		metrics.IncrCounter([]string{"nomad", "rpc", "accept_conn"}, 1)
	}
}

// handleConn is used to determine if this is a Raft or
// Nomad type RPC connection and invoke the correct handler
func (s *Server) handleConn(conn net.Conn, isTLS bool) {
	// Read a single byte
	buf := make([]byte, 1)
	if _, err := conn.Read(buf); err != nil {
		if err != io.EOF {
			s.logger.Printf("[ERR] nomad.rpc: failed to read byte: %v", err)
		}
		conn.Close()
		return
	}

	// Enforce TLS if VerifyIncoming is set
	if s.config.RequireTLS && !isTLS && RPCType(buf[0]) != rpcTLS {
		s.logger.Printf("[WARN] nomad.rpc: Non-TLS connection attempted with RequireTLS set")
		conn.Close()
		return
	}

	// Switch on the byte
	switch RPCType(buf[0]) {
	case rpcNomad:
		s.handleNomadConn(conn)

	case rpcRaft:
		metrics.IncrCounter([]string{"nomad", "rpc", "raft_handoff"}, 1)
		s.raftLayer.Handoff(conn)

	case rpcMultiplex:
		s.handleMultiplex(conn)

	case rpcTLS:
		if s.rpcTLS == nil {
			s.logger.Printf("[WARN] nomad.rpc: TLS connection attempted, server not configured for TLS")
			conn.Close()
			return
		}
		conn = tls.Server(conn, s.rpcTLS)
		s.handleConn(conn, true)

	default:
		s.logger.Printf("[ERR] nomad.rpc: unrecognized RPC byte: %v", buf[0])
		conn.Close()
		return
	}
}

// handleMultiplex is used to multiplex a single incoming connection
// using the Yamux multiplexer
func (s *Server) handleMultiplex(conn net.Conn) {
	defer conn.Close()
	conf := yamux.DefaultConfig()
	conf.LogOutput = s.config.LogOutput
	server, _ := yamux.Server(conn, conf)
	for {
		sub, err := server.Accept()
		if err != nil {
			if err != io.EOF {
				s.logger.Printf("[ERR] nomad.rpc: multiplex conn accept failed: %v", err)
			}
			return
		}
		go s.handleNomadConn(sub)
	}
}

// handleNomadConn is used to service a single Nomad RPC connection
func (s *Server) handleNomadConn(conn net.Conn) {
	defer conn.Close()
	rpcCodec := msgpackrpc.NewServerCodec(conn)
	for {
		select {
		case <-s.shutdownCh:
			return
		default:
		}

		if err := s.rpcServer.ServeRequest(rpcCodec); err != nil {
			if err != io.EOF && !strings.Contains(err.Error(), "closed") {
				s.logger.Printf("[ERR] nomad.rpc: RPC error: %v (%v)", err, conn)
				metrics.IncrCounter([]string{"nomad", "rpc", "request_error"}, 1)
			}
			return
		}
		metrics.IncrCounter([]string{"nomad", "rpc", "request"}, 1)
	}
}

// raftApply is used to encode a message, run it through raft, and return
// the FSM response along with any errors
func (s *Server) raftApply(t structs.MessageType, msg interface{}) (interface{}, error) {
	buf, err := structs.Encode(t, msg)
	if err != nil {
		return nil, fmt.Errorf("Failed to encode request: %v", err)
	}

	// Warn if the command is very large
	if n := len(buf); n > raftWarnSize {
		s.logger.Printf("[WARN] nomad: Attempting to apply large raft entry (type %d) (%d bytes)", t, n)
	}

	future := s.raft.Apply(buf, enqueueLimit)
	if err := future.Error(); err != nil {
		return nil, err
	}

	return future.Response(), nil
}
<|endoftext|>


<commit_before>package libvirt

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"os"
	"path"
	"time"
)

const (
	maxIfaceNum = 100
)

// randomMACAddress returns a randomized MAC address
// with libvirt prefix
func randomMACAddress() (string, error) {
	buf := make([]byte, 3)
	_, err := rand.Read(buf)
	if err != nil {
		return "", err
	}

	// set local bit and unicast
	buf[0] = (buf[0] | 2) & 0xfe
	// Set the local bit
	buf[0] |= 2

	// avoid libvirt-reserved addresses
	if buf[0] == 0xfe {
		buf[0] = 0xee
	}

	return fmt.Sprintf("52:54:00:%02x:%02x:%02x",
		buf[0], buf[1], buf[2]), nil
}

// randomPort returns a random port
func randomPort() int {
	const minPort = 1024
	const maxPort = 65535

	rand.Seed(time.Now().UnixNano())
	return rand.Intn(maxPort-minPort) + minPort
}

// freeNetworkInterface returns a free network interface
func freeNetworkInterface(basename string) (string, error) {
	for i := 0; i < maxIfaceNum; i++ {
		ifaceName := fmt.Sprintf("%s%d", basename, i)
		_, err := net.InterfaceByName(ifaceName)
		if err != nil {
			return ifaceName, nil
		}
	}
	return "", fmt.Errorf("could not obtain a free network interface")
}

// networkRange calculates the first and last IP addresses in an IPNet
func networkRange(network *net.IPNet) (net.IP, net.IP) {
	netIP := network.IP.To4()
	lastIP := net.IPv4(0, 0, 0, 0).To4()
	if netIP == nil {
		netIP = network.IP.To16()
		lastIP = net.IPv6zero.To16()
	}
	firstIP := netIP.Mask(network.Mask)
	for i := 0; i < len(lastIP); i++ {
		lastIP[i] = netIP[i] | ^network.Mask[i]
	}
	return firstIP, lastIP
}

// a HTTP server that serves files in a directory, used mostly for testing
type fileWebServer struct {
	Dir  string
	Port int
	URL  string

	server *http.Server
}

func (fws *fileWebServer) Start() error {
	dir, err := ioutil.TempDir(fws.Dir, "")
	if err != nil {
		return err
	}

	fws.Dir = dir
	fws.Port = randomPort()
	fws.URL = fmt.Sprintf("http://127.0.0.1:%d", fws.Port)

	handler := http.NewServeMux()
	handler.Handle("/", http.FileServer(http.Dir(dir)))
	fws.server = &http.Server{Addr: fmt.Sprintf(":%d", fws.Port), Handler: handler}
	ln, err := net.Listen("tcp", fws.server.Addr)
	if err != nil {
		return err
	}
	go fws.server.Serve(ln)
	return nil
}

// Adds a file (with some content) in the directory served by the fileWebServer
func (fws *fileWebServer) AddContent(content []byte) (string, *os.File, error) {
	tmpfile, err := ioutil.TempFile(fws.Dir, "file-")
	if err != nil {
		return "", nil, err
	}

	if len(content) > 0 {
		if _, err := tmpfile.Write(content); err != nil {
			return "", nil, err
		}
	}

	return fmt.Sprintf("%s/%s", fws.URL, path.Base(tmpfile.Name())), tmpfile, nil
}

// Symlinks a file into the directory server by the webserver
func (fws *fileWebServer) AddFile(filePath string) (string, error) {
	err := os.Symlink(filePath, path.Join(fws.Dir, path.Base(filePath)))
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("%s/%s", fws.URL, path.Base(filePath)), nil
}

func (fws *fileWebServer) Stop() {
	os.RemoveAll(fws.Dir)
}
<commit_msg>Seed the random number generator<commit_after>package libvirt

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"os"
	"path"
	"time"
)

const (
	maxIfaceNum = 100
)

// randomMACAddress returns a randomized MAC address
// with libvirt prefix
func randomMACAddress() (string, error) {
	buf := make([]byte, 3)
	rand.Seed(time.Now().UnixNano())
	_, err := rand.Read(buf)
	if err != nil {
		return "", err
	}

	// set local bit and unicast
	buf[0] = (buf[0] | 2) & 0xfe
	// Set the local bit
	buf[0] |= 2

	// avoid libvirt-reserved addresses
	if buf[0] == 0xfe {
		buf[0] = 0xee
	}

	return fmt.Sprintf("52:54:00:%02x:%02x:%02x",
		buf[0], buf[1], buf[2]), nil
}

// randomPort returns a random port
func randomPort() int {
	const minPort = 1024
	const maxPort = 65535

	rand.Seed(time.Now().UnixNano())
	return rand.Intn(maxPort-minPort) + minPort
}

// freeNetworkInterface returns a free network interface
func freeNetworkInterface(basename string) (string, error) {
	for i := 0; i < maxIfaceNum; i++ {
		ifaceName := fmt.Sprintf("%s%d", basename, i)
		_, err := net.InterfaceByName(ifaceName)
		if err != nil {
			return ifaceName, nil
		}
	}
	return "", fmt.Errorf("could not obtain a free network interface")
}

// networkRange calculates the first and last IP addresses in an IPNet
func networkRange(network *net.IPNet) (net.IP, net.IP) {
	netIP := network.IP.To4()
	lastIP := net.IPv4(0, 0, 0, 0).To4()
	if netIP == nil {
		netIP = network.IP.To16()
		lastIP = net.IPv6zero.To16()
	}
	firstIP := netIP.Mask(network.Mask)
	for i := 0; i < len(lastIP); i++ {
		lastIP[i] = netIP[i] | ^network.Mask[i]
	}
	return firstIP, lastIP
}

// a HTTP server that serves files in a directory, used mostly for testing
type fileWebServer struct {
	Dir  string
	Port int
	URL  string

	server *http.Server
}

func (fws *fileWebServer) Start() error {
	dir, err := ioutil.TempDir(fws.Dir, "")
	if err != nil {
		return err
	}

	fws.Dir = dir
	fws.Port = randomPort()
	fws.URL = fmt.Sprintf("http://127.0.0.1:%d", fws.Port)

	handler := http.NewServeMux()
	handler.Handle("/", http.FileServer(http.Dir(dir)))
	fws.server = &http.Server{Addr: fmt.Sprintf(":%d", fws.Port), Handler: handler}
	ln, err := net.Listen("tcp", fws.server.Addr)
	if err != nil {
		return err
	}
	go fws.server.Serve(ln)
	return nil
}

// Adds a file (with some content) in the directory served by the fileWebServer
func (fws *fileWebServer) AddContent(content []byte) (string, *os.File, error) {
	tmpfile, err := ioutil.TempFile(fws.Dir, "file-")
	if err != nil {
		return "", nil, err
	}

	if len(content) > 0 {
		if _, err := tmpfile.Write(content); err != nil {
			return "", nil, err
		}
	}

	return fmt.Sprintf("%s/%s", fws.URL, path.Base(tmpfile.Name())), tmpfile, nil
}

// Symlinks a file into the directory server by the webserver
func (fws *fileWebServer) AddFile(filePath string) (string, error) {
	err := os.Symlink(filePath, path.Join(fws.Dir, path.Base(filePath)))
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("%s/%s", fws.URL, path.Base(filePath)), nil
}

func (fws *fileWebServer) Stop() {
	os.RemoveAll(fws.Dir)
}
<|endoftext|>


<commit_before>package spider

import (
	"math"
	"runtime"
	"sync"
	"time"

	"github.com/henrylee2cn/pholcus/app/downloader/request"
	"github.com/henrylee2cn/pholcus/app/scheduler"
	"github.com/henrylee2cn/pholcus/common/util"
	"github.com/henrylee2cn/pholcus/logs"
	"github.com/henrylee2cn/pholcus/runtime/status"
)

const (
	KEYWORD     = util.USE_KEYWORD // 若使用Keyword，则Keyword初始值必须为USE_KEYWORD
	MAXPAGE     = math.MaxInt64    // 如希望在规则中自定义控制MaxPage，则MaxPage初始值必须为MAXPAGE
	ACTIVE_STOP = "——主动终止Spider——"
)

// 蜘蛛规则
type (
	Spider struct {
		Id   int    // 所在SpiderList的下标编号，系统自动分配
		Name string // 必须保证全局唯一
		*RuleTree

		//以下为可选成员
		Description  string
		Pausetime    int64  // 暂停区间(随机: Pausetime/2 ~ Pausetime*2)
		EnableCookie bool   // 控制所有请求是否使用cookie记录
		MaxPage      int64  // 为负值时自动在调度中限制请求数，为正值时在规则中自定义控制
		Keyword      string // 如需使用必须附初始值为常量USE_KEYWORD
		// 命名空间相对于数据库名，不依赖具体数据内容
		Namespace func(*Spider) string
		// 子命名空间相对于表名，可依赖具体数据内容
		SubNamespace func(self *Spider, dataCell map[string]interface{}) string

		// 请求矩阵
		ReqMatrix *scheduler.Matrix
		// 定时器
		timer *Timer
		// 执行状态
		status  int
		rwMutex sync.RWMutex
	}

	//采集规则树
	RuleTree struct {
		// 执行入口（树根）
		Root func(*Context)
		// 执行解析过程（树干）
		Trunk map[string]*Rule
	}

	// 采集规则单元
	Rule struct {
		// 输出结果的字段名列表
		ItemFields []string
		// 内容解析函数
		ParseFunc func(*Context)
		// 通用辅助函数
		AidFunc func(*Context, map[string]interface{}) interface{}
	}
)

// 添加自身到蜘蛛菜单
func (self *Spider) Register() *Spider {
	Menu.Add(self)
	return self
}

// 指定规则的获取结果的字段名列表
func (self *Spider) GetItemFields(rule *Rule) []string {
	return rule.ItemFields
}

// 返回结果字段名的值
// 不存在时返回空字符串
func (self *Spider) GetItemField(rule *Rule, index int) (field string) {
	if index > len(rule.ItemFields)-1 || index < 0 {
		return ""
	}
	return rule.ItemFields[index]
}

// 返回结果字段名的其索引
// 不存在时索引为-1
func (self *Spider) GetItemFieldIndex(rule *Rule, field string) (index int) {
	for idx, v := range rule.ItemFields {
		if v == field {
			return idx
		}
	}
	return -1
}

// 为指定Rule动态追加结果字段名，并返回索引位置
// 已存在时返回原来索引位置
func (self *Spider) UpsertItemField(rule *Rule, field string) (index int) {
	for i, v := range rule.ItemFields {
		if v == field {
			return i
		}
	}
	rule.ItemFields = append(rule.ItemFields, field)
	return len(rule.ItemFields) - 1
}

// 获取蜘蛛名称
func (self *Spider) GetName() string {
	return self.Name
}

// 安全返回指定规则
func (self *Spider) GetRule(ruleName string) (*Rule, bool) {
	rule, found := self.RuleTree.Trunk[ruleName]
	return rule, found
}

// 返回指定规则
func (self *Spider) MustGetRule(ruleName string) *Rule {
	return self.RuleTree.Trunk[ruleName]
}

// 返回规则树
func (self *Spider) GetRules() map[string]*Rule {
	return self.RuleTree.Trunk
}

// 获取蜘蛛描述
func (self *Spider) GetDescription() string {
	return self.Description
}

// 获取蜘蛛ID
func (self *Spider) GetId() int {
	return self.Id
}

// 设置蜘蛛ID
func (self *Spider) SetId(id int) {
	self.Id = id
}

// 获取自定义输入
func (self *Spider) GetKeyword() string {
	return self.Keyword
}

// 设置自定义输入
func (self *Spider) SetKeyword(keyword string) {
	self.Keyword = keyword
}

// 获取采集的最大页数
func (self *Spider) GetMaxPage() int64 {
	return self.MaxPage
}

// 设置采集的最大页数
func (self *Spider) SetMaxPage(max int64) {
	self.MaxPage = max
}

// 控制所有请求是否使用cookie
func (self *Spider) GetEnableCookie() bool {
	return self.EnableCookie
}

// 自定义暂停时间 pause[0]~(pause[0]+pause[1])，优先级高于外部传参
// 当且仅当runtime[0]为true时可覆盖现有值
func (self *Spider) SetPausetime(pause int64, runtime ...bool) {
	if self.Pausetime == 0 || len(runtime) > 0 && runtime[0] {
		self.Pausetime = pause
	}
}

// 设置定时器
// @id为定时器唯一标识
// @bell==nil时为倒计时器，此时@tol为睡眠时长
// @bell!=nil时为闹铃，此时@tol用于指定醒来时刻（从now起遇到的第tol个bell）
func (self *Spider) SetTimer(id string, tol time.Duration, bell *Bell) bool {
	if self.timer == nil {
		self.timer = newTimer()
	}
	return self.timer.set(id, tol, bell)
}

// 启动定时器，并返回定时器是否可以继续使用
func (self *Spider) RunTimer(id string) bool {
	if self.timer == nil {
		return false
	}
	return self.timer.sleep(id)
}

// 返回一个自身复制品
func (self *Spider) Copy() *Spider {
	ghost := &Spider{}
	ghost.Name = self.Name

	ghost.RuleTree = &RuleTree{
		Root:  self.Root,
		Trunk: make(map[string]*Rule, len(self.RuleTree.Trunk)),
	}
	for k, v := range self.RuleTree.Trunk {
		ghost.RuleTree.Trunk[k] = new(Rule)

		ghost.RuleTree.Trunk[k].ItemFields = make([]string, len(v.ItemFields))
		copy(ghost.RuleTree.Trunk[k].ItemFields, v.ItemFields)

		ghost.RuleTree.Trunk[k].ParseFunc = v.ParseFunc
		ghost.RuleTree.Trunk[k].AidFunc = v.AidFunc
	}

	ghost.Description = self.Description
	ghost.Pausetime = self.Pausetime
	ghost.EnableCookie = self.EnableCookie
	ghost.MaxPage = self.MaxPage
	ghost.Keyword = self.Keyword

	ghost.Namespace = self.Namespace
	ghost.SubNamespace = self.SubNamespace

	ghost.timer = self.timer
	ghost.status = self.status

	return ghost
}

func (self *Spider) ReqmatrixInit() *Spider {
	if self.MaxPage < 0 {
		self.ReqMatrix = scheduler.AddMatrix(self.GetName(), self.MaxPage)
		self.SetMaxPage(0)
	} else {
		self.ReqMatrix = scheduler.AddMatrix(self.GetName(), math.MinInt64)
	}
	return self
}

// 返回是否作为新的失败请求被添加至队列尾部
func (self *Spider) DoHistory(req *request.Request, ok bool) bool {
	return self.ReqMatrix.DoHistory(req, ok)
}

func (self *Spider) RequestPush(req *request.Request) {
	self.ReqMatrix.Push(req)
}

func (self *Spider) RequestPull() *request.Request {
	return self.ReqMatrix.Pull()
}

func (self *Spider) RequestUse() {
	self.ReqMatrix.Use()
}

func (self *Spider) RequestFree() {
	self.ReqMatrix.Free()
}

func (self *Spider) CanStop() bool {
	return self.ReqMatrix.CanStop()
}

func (self *Spider) RequestLen() int {
	return self.ReqMatrix.Len()
}

func (self *Spider) TryFlushHistory() {
	self.ReqMatrix.TryFlushHistory()
}

// 开始执行蜘蛛
func (self *Spider) Start() {
	self.rwMutex.Lock()
	self.status = status.RUN
	self.rwMutex.Unlock()
	defer func() {
		if p := recover(); p != nil {
			logs.Log.Error(" *     Panic  [root]: %v\n", p)
		}
	}()
	self.RuleTree.Root(NewContext(self, nil))
	cancel := time.After(1e9)
	for self.RequestLen() == 0 {
		select {
		case <-cancel:
			return
		default:
			runtime.Gosched()
		}
	}
}

// 主动崩溃爬虫运行协程
func (self *Spider) Stop() {
	self.rwMutex.Lock()
	self.status = status.STOP
	// 取消所有定时器
	if self.timer != nil {
		self.timer.drop()
		self.timer = nil
	}
	self.rwMutex.Unlock()
}

// 若已主动终止任务，则崩溃爬虫协程
func (self *Spider) tryPanic() {
	self.rwMutex.RLock()
	if self.status == status.STOP {
		self.rwMutex.RUnlock()
		panic(ACTIVE_STOP)
	}
	self.rwMutex.RUnlock()
}

// 退出任务前收尾工作
func (self *Spider) Defer() {
	// 取消所有定时器
	if self.timer != nil {
		self.timer.drop()
		self.timer = nil
	}
	// 等待处理中的请求完成
	self.ReqMatrix.Wait()
}
<commit_msg>允许在规则的Root()中进行阻塞<commit_after>package spider

import (
	"math"
	"runtime"
	"sync"
	"time"

	"github.com/henrylee2cn/pholcus/app/downloader/request"
	"github.com/henrylee2cn/pholcus/app/scheduler"
	"github.com/henrylee2cn/pholcus/common/util"
	"github.com/henrylee2cn/pholcus/logs"
	"github.com/henrylee2cn/pholcus/runtime/status"
)

const (
	KEYWORD     = util.USE_KEYWORD // 若使用Keyword，则Keyword初始值必须为USE_KEYWORD
	MAXPAGE     = math.MaxInt64    // 如希望在规则中自定义控制MaxPage，则MaxPage初始值必须为MAXPAGE
	ACTIVE_STOP = "——主动终止Spider——"
)

// 蜘蛛规则
type (
	Spider struct {
		Id   int    // 所在SpiderList的下标编号，系统自动分配
		Name string // 必须保证全局唯一
		*RuleTree

		//以下为可选成员
		Description  string
		Pausetime    int64  // 暂停区间(随机: Pausetime/2 ~ Pausetime*2)
		EnableCookie bool   // 控制所有请求是否使用cookie记录
		MaxPage      int64  // 为负值时自动在调度中限制请求数，为正值时在规则中自定义控制
		Keyword      string // 如需使用必须附初始值为常量USE_KEYWORD
		// 命名空间相对于数据库名，不依赖具体数据内容
		Namespace func(*Spider) string
		// 子命名空间相对于表名，可依赖具体数据内容
		SubNamespace func(self *Spider, dataCell map[string]interface{}) string

		// 请求矩阵
		ReqMatrix *scheduler.Matrix
		// 定时器
		timer *Timer
		// 执行状态
		status  int
		rwMutex sync.RWMutex
	}

	//采集规则树
	RuleTree struct {
		// 执行入口（树根）
		Root func(*Context)
		// 执行解析过程（树干）
		Trunk map[string]*Rule
	}

	// 采集规则单元
	Rule struct {
		// 输出结果的字段名列表
		ItemFields []string
		// 内容解析函数
		ParseFunc func(*Context)
		// 通用辅助函数
		AidFunc func(*Context, map[string]interface{}) interface{}
	}
)

// 添加自身到蜘蛛菜单
func (self *Spider) Register() *Spider {
	Menu.Add(self)
	return self
}

// 指定规则的获取结果的字段名列表
func (self *Spider) GetItemFields(rule *Rule) []string {
	return rule.ItemFields
}

// 返回结果字段名的值
// 不存在时返回空字符串
func (self *Spider) GetItemField(rule *Rule, index int) (field string) {
	if index > len(rule.ItemFields)-1 || index < 0 {
		return ""
	}
	return rule.ItemFields[index]
}

// 返回结果字段名的其索引
// 不存在时索引为-1
func (self *Spider) GetItemFieldIndex(rule *Rule, field string) (index int) {
	for idx, v := range rule.ItemFields {
		if v == field {
			return idx
		}
	}
	return -1
}

// 为指定Rule动态追加结果字段名，并返回索引位置
// 已存在时返回原来索引位置
func (self *Spider) UpsertItemField(rule *Rule, field string) (index int) {
	for i, v := range rule.ItemFields {
		if v == field {
			return i
		}
	}
	rule.ItemFields = append(rule.ItemFields, field)
	return len(rule.ItemFields) - 1
}

// 获取蜘蛛名称
func (self *Spider) GetName() string {
	return self.Name
}

// 安全返回指定规则
func (self *Spider) GetRule(ruleName string) (*Rule, bool) {
	rule, found := self.RuleTree.Trunk[ruleName]
	return rule, found
}

// 返回指定规则
func (self *Spider) MustGetRule(ruleName string) *Rule {
	return self.RuleTree.Trunk[ruleName]
}

// 返回规则树
func (self *Spider) GetRules() map[string]*Rule {
	return self.RuleTree.Trunk
}

// 获取蜘蛛描述
func (self *Spider) GetDescription() string {
	return self.Description
}

// 获取蜘蛛ID
func (self *Spider) GetId() int {
	return self.Id
}

// 设置蜘蛛ID
func (self *Spider) SetId(id int) {
	self.Id = id
}

// 获取自定义输入
func (self *Spider) GetKeyword() string {
	return self.Keyword
}

// 设置自定义输入
func (self *Spider) SetKeyword(keyword string) {
	self.Keyword = keyword
}

// 获取采集的最大页数
func (self *Spider) GetMaxPage() int64 {
	return self.MaxPage
}

// 设置采集的最大页数
func (self *Spider) SetMaxPage(max int64) {
	self.MaxPage = max
}

// 控制所有请求是否使用cookie
func (self *Spider) GetEnableCookie() bool {
	return self.EnableCookie
}

// 自定义暂停时间 pause[0]~(pause[0]+pause[1])，优先级高于外部传参
// 当且仅当runtime[0]为true时可覆盖现有值
func (self *Spider) SetPausetime(pause int64, runtime ...bool) {
	if self.Pausetime == 0 || len(runtime) > 0 && runtime[0] {
		self.Pausetime = pause
	}
}

// 设置定时器
// @id为定时器唯一标识
// @bell==nil时为倒计时器，此时@tol为睡眠时长
// @bell!=nil时为闹铃，此时@tol用于指定醒来时刻（从now起遇到的第tol个bell）
func (self *Spider) SetTimer(id string, tol time.Duration, bell *Bell) bool {
	if self.timer == nil {
		self.timer = newTimer()
	}
	return self.timer.set(id, tol, bell)
}

// 启动定时器，并返回定时器是否可以继续使用
func (self *Spider) RunTimer(id string) bool {
	if self.timer == nil {
		return false
	}
	return self.timer.sleep(id)
}

// 返回一个自身复制品
func (self *Spider) Copy() *Spider {
	ghost := &Spider{}
	ghost.Name = self.Name

	ghost.RuleTree = &RuleTree{
		Root:  self.Root,
		Trunk: make(map[string]*Rule, len(self.RuleTree.Trunk)),
	}
	for k, v := range self.RuleTree.Trunk {
		ghost.RuleTree.Trunk[k] = new(Rule)

		ghost.RuleTree.Trunk[k].ItemFields = make([]string, len(v.ItemFields))
		copy(ghost.RuleTree.Trunk[k].ItemFields, v.ItemFields)

		ghost.RuleTree.Trunk[k].ParseFunc = v.ParseFunc
		ghost.RuleTree.Trunk[k].AidFunc = v.AidFunc
	}

	ghost.Description = self.Description
	ghost.Pausetime = self.Pausetime
	ghost.EnableCookie = self.EnableCookie
	ghost.MaxPage = self.MaxPage
	ghost.Keyword = self.Keyword

	ghost.Namespace = self.Namespace
	ghost.SubNamespace = self.SubNamespace

	ghost.timer = self.timer
	ghost.status = self.status

	return ghost
}

func (self *Spider) ReqmatrixInit() *Spider {
	if self.MaxPage < 0 {
		self.ReqMatrix = scheduler.AddMatrix(self.GetName(), self.MaxPage)
		self.SetMaxPage(0)
	} else {
		self.ReqMatrix = scheduler.AddMatrix(self.GetName(), math.MinInt64)
	}
	return self
}

// 返回是否作为新的失败请求被添加至队列尾部
func (self *Spider) DoHistory(req *request.Request, ok bool) bool {
	return self.ReqMatrix.DoHistory(req, ok)
}

func (self *Spider) RequestPush(req *request.Request) {
	self.ReqMatrix.Push(req)
}

func (self *Spider) RequestPull() *request.Request {
	return self.ReqMatrix.Pull()
}

func (self *Spider) RequestUse() {
	self.ReqMatrix.Use()
}

func (self *Spider) RequestFree() {
	self.ReqMatrix.Free()
}

func (self *Spider) CanStop() bool {
	return self.ReqMatrix.CanStop()
}

func (self *Spider) RequestLen() int {
	return self.ReqMatrix.Len()
}

func (self *Spider) TryFlushHistory() {
	self.ReqMatrix.TryFlushHistory()
}

// 开始执行蜘蛛
func (self *Spider) Start() {
	self.rwMutex.Lock()
	self.status = status.RUN
	self.rwMutex.Unlock()
	defer func() {
		if p := recover(); p != nil {
			logs.Log.Error(" *     Panic  [root]: %v\n", p)
		}
	}()
	go self.RuleTree.Root(NewContext(self, nil))
	cancel := time.After(1e9)
	for self.RequestLen() == 0 {
		select {
		case <-cancel:
			return
		default:
			runtime.Gosched()
		}
	}
}

// 主动崩溃爬虫运行协程
func (self *Spider) Stop() {
	self.rwMutex.Lock()
	self.status = status.STOP
	// 取消所有定时器
	if self.timer != nil {
		self.timer.drop()
		self.timer = nil
	}
	self.rwMutex.Unlock()
}

// 若已主动终止任务，则崩溃爬虫协程
func (self *Spider) tryPanic() {
	self.rwMutex.RLock()
	if self.status == status.STOP {
		self.rwMutex.RUnlock()
		panic(ACTIVE_STOP)
	}
	self.rwMutex.RUnlock()
}

// 退出任务前收尾工作
func (self *Spider) Defer() {
	// 取消所有定时器
	if self.timer != nil {
		self.timer.drop()
		self.timer = nil
	}
	// 等待处理中的请求完成
	self.ReqMatrix.Wait()
}
<|endoftext|>


<commit_before>package litrpc

import (
	"fmt"

	"github.com/adiabat/bech32"
	"github.com/adiabat/btcd/chaincfg"
	"github.com/adiabat/btcd/txscript"
	"github.com/adiabat/btcd/wire"
	"github.com/adiabat/btcutil"
	"github.com/mit-dci/lit/portxo"
)

type TxidsReply struct {
	Txids []string
}
type StatusReply struct {
	Status string
}

type NoArgs struct {
	// nothin
}

// ------------------------- balance
// BalReply is the reply when the user asks about their balance.
// This is a Non-Channel
type BalReply struct {
	ChanTotal   int64 // total balance in channels
	TxoTotal    int64 // all utxos
	MatureWitty int64 // confirmed, spendable and witness
}

func (r *LitRPC) Bal(args *NoArgs, reply *BalReply) error {
	var err error
	var allTxos portxo.TxoSliceByAmt

	nowHeight := r.Node.SubWallet.CurrentHeight()

	allTxos, err = r.Node.SubWallet.UtxoDump()
	if err != nil {
		return err
	}

	// ask sub-wallet for balance
	reply.TxoTotal = allTxos.Sum()
	reply.MatureWitty = allTxos.SumWitness(nowHeight)

	// get all channel states
	qcs, err := r.Node.GetAllQchans()
	if err != nil {
		return err
	}
	// iterate through channels to figure out how much we have
	for _, q := range qcs {
		reply.ChanTotal += q.State.MyAmt
	}

	return nil
}

type TxoInfo struct {
	OutPoint string
	Amt      int64
	Height   int32
	Delay    int32
	Witty    bool

	KeyPath string
}
type TxoListReply struct {
	Txos []TxoInfo
}

// TxoList sends back a list of all non-channel utxos
func (r *LitRPC) TxoList(args *NoArgs, reply *TxoListReply) error {
	allTxos, err := r.Node.SubWallet.UtxoDump()
	if err != nil {
		return err
	}

	syncHeight := r.Node.SubWallet.CurrentHeight()

	reply.Txos = make([]TxoInfo, len(allTxos))
	for i, u := range allTxos {
		reply.Txos[i].OutPoint = u.Op.String()
		reply.Txos[i].Amt = u.Value
		reply.Txos[i].Height = u.Height
		// show delay before utxo can be spent
		if u.Seq != 0 {
			reply.Txos[i].Delay = u.Height + int32(u.Seq) - syncHeight
		}
		reply.Txos[i].Witty = u.Mode&portxo.FlagTxoWitness != 0
		reply.Txos[i].KeyPath = u.KeyGen.String()
	}
	return nil
}

type SyncHeightReply struct {
	SyncHeight   int32
	HeaderHeight int32
}

func (r *LitRPC) SyncHeight(args *NoArgs, reply *SyncHeightReply) error {
	reply.SyncHeight = r.Node.SubWallet.CurrentHeight()
	return nil
}

// ------------------------- send
type SendArgs struct {
	DestAddrs []string
	Amts      []int64
}

func (r *LitRPC) Send(args SendArgs, reply *TxidsReply) error {
	var err error

	nOutputs := len(args.DestAddrs)
	if nOutputs < 1 {
		return fmt.Errorf("No destination address specified")
	}
	if nOutputs != len(args.Amts) {
		return fmt.Errorf("%d addresses but %d amounts specified",
			nOutputs, len(args.Amts))
	}

	txOuts := make([]*wire.TxOut, nOutputs)
	for i, s := range args.DestAddrs {
		if args.Amts[i] < 10000 {
			return fmt.Errorf("Amt %d less than min 10000", args.Amts[i])
		}

		outScript, err := AdrStringToOutscript(s, r.Node.SubWallet.Params())
		if err != nil {
			return err
		}

		txOuts[i] = wire.NewTxOut(args.Amts[i], outScript)
	}

	// we don't care if it's witness or not
	ops, err := r.Node.SubWallet.MaybeSend(txOuts, false)
	if err != nil {
		return err
	}

	err = r.Node.SubWallet.ReallySend(&ops[0].Hash)
	if err != nil {
		return err
	}

	reply.Txids = append(reply.Txids, ops[0].Hash.String())
	return nil
}

// ------------------------- sweep
type SweepArgs struct {
	DestAdr string
	NumTx   uint32
	Drop    bool
}

// AdrStringToOutscript converts an address string into an output script byte slice
func AdrStringToOutscript(adr string, p *chaincfg.Params) ([]byte, error) {
	var err error
	var outScript []byte

	// try bech32 address
	_, outScript, err = bech32.Decode(adr)
	if err == nil {
		return outScript, nil
	}

	// try for base58 address
	oadr, err := btcutil.DecodeAddress(adr, p)
	if err != nil {
		return nil, err
	}
	outScript, err = txscript.PayToAddrScript(oadr)
	if err != nil {
		return nil, err
	}
	return outScript, nil
}

func (r *LitRPC) Sweep(args SweepArgs, reply *TxidsReply) error {

	outScript, err := AdrStringToOutscript(args.DestAdr, r.Node.SubWallet.Params())
	if err != nil {
		return err
	}

	fmt.Printf("numtx: %d\n", args.NumTx)
	if args.NumTx < 1 {
		return fmt.Errorf("can't send %d txs", args.NumTx)
	}

	txids, err := r.Node.SubWallet.Sweep(outScript, args.NumTx)
	if err != nil {
		return err
	}

	for _, txid := range txids {
		reply.Txids = append(reply.Txids, txid.String())
	}

	return nil
}

// ------------------------- fanout
type FanArgs struct {
	DestAdr      string
	NumOutputs   uint32
	AmtPerOutput int64
}

func (r *LitRPC) Fanout(args FanArgs, reply *TxidsReply) error {
	if args.NumOutputs < 1 {
		return fmt.Errorf("Must have at least 1 output")
	}
	if args.AmtPerOutput < 5000 {
		return fmt.Errorf("Minimum 5000 per output")
	}
	outScript, err := AdrStringToOutscript(args.DestAdr, r.Node.SubWallet.Params())
	if err != nil {
		return err
	}

	txos := make([]*wire.TxOut, args.NumOutputs)

	for i, _ := range txos {
		txos[i] = new(wire.TxOut)
		txos[i].Value = args.AmtPerOutput + int64(i)
		txos[i].PkScript = outScript
	}

	// don't care if inputs are witty or not
	ops, err := r.Node.SubWallet.MaybeSend(txos, false)
	if err != nil {
		return err
	}
	err = r.Node.SubWallet.ReallySend(&ops[0].Hash)
	if err != nil {
		return err
	}

	reply.Txids = append(reply.Txids, ops[0].String())
	return nil
}

// ------------------------- address
type AddressArgs struct {
	NumToMake uint32
}
type AddressReply struct {
	WitAddresses    []string
	LegacyAddresses []string
}

func (r *LitRPC) Address(args *AddressArgs, reply *AddressReply) error {
	var err error
	var allAdr [][20]byte

	// If you tell it to make 0 new addresses, it sends a list of all the old ones
	if args.NumToMake == 0 {
		// this gets 20 byte addresses; need to convert them to bech32 / base58
		allAdr, err = r.Node.SubWallet.AdrDump()
		if err != nil {
			return err
		}
	} else {
		// call NewAdr a bunch of times
		remaining := args.NumToMake
		for remaining > 0 {
			adr, err := r.Node.SubWallet.NewAdr()
			if err != nil {
				return err
			}
			allAdr = append(allAdr, adr)
			remaining--
		}
	}
	reply.WitAddresses = make([]string, len(allAdr))
	reply.LegacyAddresses = make([]string, len(allAdr))

	for i, a := range allAdr {
		// convert 20 byte array to old address
		oldadr, err := btcutil.NewAddressPubKeyHash(
			a[:], r.Node.SubWallet.Params())
		if err != nil {
			return err
		}
		reply.LegacyAddresses[i] = oldadr.String()

		// convert 20-byte PKH to a bech32 segwit v0 address
		bech32adr := bech32.Encode(r.Node.SubWallet.Params().Bech32Prefix, a[:])
		if err != nil {
			return err
		}
		reply.WitAddresses[i] = bech32adr
	}

	return nil
}
<commit_msg>fix nonstandard bech32 addresses<commit_after>package litrpc

import (
	"fmt"

	"github.com/adiabat/bech32"
	"github.com/adiabat/btcd/chaincfg"
	"github.com/adiabat/btcd/txscript"
	"github.com/adiabat/btcd/wire"
	"github.com/adiabat/btcutil"
	"github.com/mit-dci/lit/portxo"
)

type TxidsReply struct {
	Txids []string
}
type StatusReply struct {
	Status string
}

type NoArgs struct {
	// nothin
}

// ------------------------- balance
// BalReply is the reply when the user asks about their balance.
// This is a Non-Channel
type BalReply struct {
	ChanTotal   int64 // total balance in channels
	TxoTotal    int64 // all utxos
	MatureWitty int64 // confirmed, spendable and witness
}

func (r *LitRPC) Bal(args *NoArgs, reply *BalReply) error {
	var err error
	var allTxos portxo.TxoSliceByAmt

	nowHeight := r.Node.SubWallet.CurrentHeight()

	allTxos, err = r.Node.SubWallet.UtxoDump()
	if err != nil {
		return err
	}

	// ask sub-wallet for balance
	reply.TxoTotal = allTxos.Sum()
	reply.MatureWitty = allTxos.SumWitness(nowHeight)

	// get all channel states
	qcs, err := r.Node.GetAllQchans()
	if err != nil {
		return err
	}
	// iterate through channels to figure out how much we have
	for _, q := range qcs {
		reply.ChanTotal += q.State.MyAmt
	}

	return nil
}

type TxoInfo struct {
	OutPoint string
	Amt      int64
	Height   int32
	Delay    int32
	Witty    bool

	KeyPath string
}
type TxoListReply struct {
	Txos []TxoInfo
}

// TxoList sends back a list of all non-channel utxos
func (r *LitRPC) TxoList(args *NoArgs, reply *TxoListReply) error {
	allTxos, err := r.Node.SubWallet.UtxoDump()
	if err != nil {
		return err
	}

	syncHeight := r.Node.SubWallet.CurrentHeight()

	reply.Txos = make([]TxoInfo, len(allTxos))
	for i, u := range allTxos {
		reply.Txos[i].OutPoint = u.Op.String()
		reply.Txos[i].Amt = u.Value
		reply.Txos[i].Height = u.Height
		// show delay before utxo can be spent
		if u.Seq != 0 {
			reply.Txos[i].Delay = u.Height + int32(u.Seq) - syncHeight
		}
		reply.Txos[i].Witty = u.Mode&portxo.FlagTxoWitness != 0
		reply.Txos[i].KeyPath = u.KeyGen.String()
	}
	return nil
}

type SyncHeightReply struct {
	SyncHeight   int32
	HeaderHeight int32
}

func (r *LitRPC) SyncHeight(args *NoArgs, reply *SyncHeightReply) error {
	reply.SyncHeight = r.Node.SubWallet.CurrentHeight()
	return nil
}

// ------------------------- send
type SendArgs struct {
	DestAddrs []string
	Amts      []int64
}

func (r *LitRPC) Send(args SendArgs, reply *TxidsReply) error {
	var err error

	nOutputs := len(args.DestAddrs)
	if nOutputs < 1 {
		return fmt.Errorf("No destination address specified")
	}
	if nOutputs != len(args.Amts) {
		return fmt.Errorf("%d addresses but %d amounts specified",
			nOutputs, len(args.Amts))
	}

	txOuts := make([]*wire.TxOut, nOutputs)
	for i, s := range args.DestAddrs {
		if args.Amts[i] < 10000 {
			return fmt.Errorf("Amt %d less than min 10000", args.Amts[i])
		}

		outScript, err := AdrStringToOutscript(s)
		if err != nil {
			return err
		}

		txOuts[i] = wire.NewTxOut(args.Amts[i], outScript)
	}

	// we don't care if it's witness or not
	ops, err := r.Node.SubWallet.MaybeSend(txOuts, false)
	if err != nil {
		return err
	}

	err = r.Node.SubWallet.ReallySend(&ops[0].Hash)
	if err != nil {
		return err
	}

	reply.Txids = append(reply.Txids, ops[0].Hash.String())
	return nil
}

// ------------------------- sweep
type SweepArgs struct {
	DestAdr string
	NumTx   uint32
	Drop    bool
}

// AdrStringToOutscript converts an address string into an output script byte slice
func AdrStringToOutscript(adr string) ([]byte, error) {
	var err error
	var outScript []byte

	// use HRP to determine network / wallet to use
	outScript, err = bech32.SegWitAddressDecode(adr)
	if err != nil { // valid bech32 string
		// try for base58 address
		// btcutil addresses don't really work as they won't tell you the
		// network; you have to tell THEM the network, which defeats the point
		// of having an address.  default to testnet only here

		// could work on adding more old-style addresses; for now use new bech32
		// addresses for multi-wallet / segwit sends.
		adr, err := btcutil.DecodeAddress(adr, &chaincfg.TestNet3Params)
		if err != nil {
			return nil, err
		}

		outScript, err = txscript.PayToAddrScript(adr)
		if err != nil {
			return nil, err
		}
	}
	return outScript, nil
}

func (r *LitRPC) Sweep(args SweepArgs, reply *TxidsReply) error {

	outScript, err := AdrStringToOutscript(args.DestAdr)
	if err != nil {
		return err
	}

	fmt.Printf("numtx: %d\n", args.NumTx)
	if args.NumTx < 1 {
		return fmt.Errorf("can't send %d txs", args.NumTx)
	}

	txids, err := r.Node.SubWallet.Sweep(outScript, args.NumTx)
	if err != nil {
		return err
	}

	for _, txid := range txids {
		reply.Txids = append(reply.Txids, txid.String())
	}

	return nil
}

// ------------------------- fanout
type FanArgs struct {
	DestAdr      string
	NumOutputs   uint32
	AmtPerOutput int64
}

func (r *LitRPC) Fanout(args FanArgs, reply *TxidsReply) error {
	if args.NumOutputs < 1 {
		return fmt.Errorf("Must have at least 1 output")
	}
	if args.AmtPerOutput < 5000 {
		return fmt.Errorf("Minimum 5000 per output")
	}
	outScript, err := AdrStringToOutscript(args.DestAdr)
	if err != nil {
		return err
	}

	txos := make([]*wire.TxOut, args.NumOutputs)

	for i, _ := range txos {
		txos[i] = new(wire.TxOut)
		txos[i].Value = args.AmtPerOutput + int64(i)
		txos[i].PkScript = outScript
	}

	// don't care if inputs are witty or not
	ops, err := r.Node.SubWallet.MaybeSend(txos, false)
	if err != nil {
		return err
	}
	err = r.Node.SubWallet.ReallySend(&ops[0].Hash)
	if err != nil {
		return err
	}

	reply.Txids = append(reply.Txids, ops[0].String())
	return nil
}

// ------------------------- address
type AddressArgs struct {
	NumToMake uint32
}
type AddressReply struct {
	WitAddresses    []string
	LegacyAddresses []string
}

func (r *LitRPC) Address(args *AddressArgs, reply *AddressReply) error {
	var err error
	var allAdr [][20]byte

	// If you tell it to make 0 new addresses, it sends a list of all the old ones
	if args.NumToMake == 0 {
		// this gets 20 byte addresses; need to convert them to bech32 / base58
		allAdr, err = r.Node.SubWallet.AdrDump()
		if err != nil {
			return err
		}
	} else {
		// call NewAdr a bunch of times
		remaining := args.NumToMake
		for remaining > 0 {
			adr, err := r.Node.SubWallet.NewAdr()
			if err != nil {
				return err
			}
			allAdr = append(allAdr, adr)
			remaining--
		}
	}
	reply.WitAddresses = make([]string, len(allAdr))
	reply.LegacyAddresses = make([]string, len(allAdr))

	for i, a := range allAdr {
		// convert 20 byte array to old address
		oldadr, err := btcutil.NewAddressPubKeyHash(
			a[:], r.Node.SubWallet.Params())
		if err != nil {
			return err
		}
		reply.LegacyAddresses[i] = oldadr.String()

		// convert 20-byte PKH to a bech32 segwit v0 address
		bech32adr, err := bech32.SegWitV0Encode(
			r.Node.SubWallet.Params().Bech32Prefix, a[:])
		if err != nil {
			return err
		}
		reply.WitAddresses[i] = bech32adr
	}

	return nil
}
<|endoftext|>


<commit_before>package glx

// #cgo linux LDFLAGS: -lGL
// #include <stdlib.h>
// #include <GL/glx.h>
import "C"
import "unsafe"

func GetProcAddress(name string) unsafe.Pointer {
	var cname *C.GLubyte = (*C.GLubyte)(C.CString(name))
	defer C.free(unsafe.Pointer(cname))
	return unsafe.Pointer(C.glXGetProcAddress(cname))
}
<commit_msg>Fix the GLX C string conversion.<commit_after>package glx

// #cgo linux LDFLAGS: -lGL
// #include <stdlib.h>
// #include <GL/glx.h>
import "C"
import "unsafe"

func GetProcAddress(name string) unsafe.Pointer {
	var cname *C.GLubyte = (*C.GLubyte)(unsafe.Pointer(C.CString(name)))
	defer C.free(unsafe.Pointer(cname))
	return unsafe.Pointer(C.glXGetProcAddress(cname))
}
<|endoftext|>


<commit_before>package ionic

import (
	"fmt"
	"net/http"
	"testing"

	. "github.com/franela/goblin"
	"github.com/gomicro/bogus"
	. "github.com/onsi/gomega"
)

func TestSearches(t *testing.T) {
	g := Goblin(t)
	RegisterFailHandler(func(m string, _ ...int) { g.Fail(m) })
	g.Describe("Searches", func() {
		var server *bogus.Bogus
		var h, p string
		var client *IonClient
		g.BeforeEach(func() {
			server = bogus.New()
			h, p = server.HostPort()
			client, _ = New(fmt.Sprintf("http://%v:%v", h, p))
		})
		g.AfterEach(func() {
			server.Close()
		})
		g.It("should perform a productidentifiers search", func() {
			server.AddPath("/v1/search/productidentifiers").
				SetMethods("GET").
				SetPayload([]byte(sampleValidProdIDSearch)).
				SetStatus(http.StatusOK)
			searchResults, err := client.GetSearchProductIdentifiers("prodIdent", "version", "vendor", "sometoken")
			Expect(err).NotTo(HaveOccurred())
			Expect(searchResults).To(HaveLen(2))
			hitRecords := server.HitRecords()
			Expect(hitRecords).To(HaveLen(1))
			Expect(hitRecords[0].Query.Get("productidentifiers")).To(Equal("prodIdent"))
			Expect(hitRecords[0].Query.Get("version")).To(Equal("version"))
			Expect(hitRecords[0].Query.Get("vendor")).To(Equal("vendor"))
			Expect(hitRecords[0].Header.Get("Authorization")).To(Equal("Bearer sometoken"))

			Expect(searchResults[0].Product.Name).To(Equal("jdk"))
			Expect(searchResults[1].Github.CommitterCount).To(Equal(uint(31337)))
		})
		g.It("should perform a repositories search", func() {
			server.AddPath("/v1/search/repositories").
				SetMethods("GET").
				SetPayload([]byte(sampleValidRositaRepoSearch)).
				SetStatus(http.StatusOK)
			searchResults, err := client.GetSearchRepositories("https://github.com/apache/ode", "sometoken")
			Expect(err).NotTo(HaveOccurred())
			Expect(searchResults).To(HaveLen(1))
			hitRecords := server.HitRecords()
			Expect(hitRecords).To(HaveLen(1))
			Expect(hitRecords[0].Query.Get("github")).To(Equal("https://github.com/apache/ode"))
			Expect(hitRecords[0].Header.Get("Authorization")).To(Equal("Bearer sometoken"))

			Expect(searchResults[0].Product.Name).To(Equal("ode"))
			Expect(searchResults[0].Product.Title).To(Equal("Apache Ode 1.1.1"))
			Expect(searchResults[0].Github.CommitterCount).To(Equal(uint(11)))
		})
	})
}

const (
	sampleValidProdIDSearch     = `{"data": [{"product": {"name": "jdk", "language": "", "title": "IBM Jdk 1.4.2", "created_at": "2017-02-13T20:02:38.33Z", "up": "", "updated_at": "2017-02-13T20:02:38.33Z", "edition": "", "version": "1.4.2", "references": [], "source": null, "org": "ibm", "part": "/a", "external_id": "cpe:/a:ibm:jdk:1.4.2", "id": 52994, "aliases": null}, "github": {"committer_count": 31337, "uri": "https://github.com/ibm/jdk"}}, {"product": {"name": "jdk", "language": "", "title": "IBM Jdk 5.0", "created_at": "2017-02-13T20:02:38.33Z", "up": "", "updated_at": "2017-02-13T20:02:38.33Z", "edition": "", "version": "5.0", "references": [], "source": null, "org": "ibm", "part": "/a", "external_id": "cpe:/a:ibm:jdk:5.0", "id": 52995, "aliases": null}, "github": {"committer_count": 31337, "uri": "https://github.com/ibm/jdk"}}], "meta": { "authors": [ "Ion Channel Dev Team" ], "copyright": "Copyright 2018 Selection Pressure LLC www.selectpress.net", "offset": 0, "total_count": 2, "version": "v1"}}`
	sampleValidRositaRepoSearch = `{"meta": {"total_count": 2, "offset": 0, "version": "v1", "copyright": "Copyright 2018 Selection Pressure LLC www.selectpress.net", "authors": ["Ion Channel Dev Team"]}, "data": [{"product": {"name": "ode", "language": "", "source": null, "created_at": "2018-03-27T22:26:14.991Z", "title": "Apache Ode 1.1.1", "up": "rc1", "updated_at": "2018-03-27T22:26:14.991Z", "edition": "", "part": "/a", "references": null, "version": "1.1.1", "org": "apache", "external_id": "cpe:/a:apache:ode:1.1.1:rc1", "id": 759066869, "aliases": null}, "github": {"committer_count": 11, "uri": "https://github.com/apache/ode"}}]}`
)
<commit_msg>responding to code review, removing internal service names from public library<commit_after>package ionic

import (
	"fmt"
	"net/http"
	"testing"

	. "github.com/franela/goblin"
	"github.com/gomicro/bogus"
	. "github.com/onsi/gomega"
)

func TestSearches(t *testing.T) {
	g := Goblin(t)
	RegisterFailHandler(func(m string, _ ...int) { g.Fail(m) })
	g.Describe("Searches", func() {
		var server *bogus.Bogus
		var h, p string
		var client *IonClient
		g.BeforeEach(func() {
			server = bogus.New()
			h, p = server.HostPort()
			client, _ = New(fmt.Sprintf("http://%v:%v", h, p))
		})
		g.AfterEach(func() {
			server.Close()
		})
		g.It("should perform a productidentifiers search", func() {
			server.AddPath("/v1/search/productidentifiers").
				SetMethods("GET").
				SetPayload([]byte(sampleValidProdIDSearch)).
				SetStatus(http.StatusOK)
			searchResults, err := client.GetSearchProductIdentifiers("prodIdent", "version", "vendor", "sometoken")
			Expect(err).NotTo(HaveOccurred())
			Expect(searchResults).To(HaveLen(2))
			hitRecords := server.HitRecords()
			Expect(hitRecords).To(HaveLen(1))
			Expect(hitRecords[0].Query.Get("productidentifiers")).To(Equal("prodIdent"))
			Expect(hitRecords[0].Query.Get("version")).To(Equal("version"))
			Expect(hitRecords[0].Query.Get("vendor")).To(Equal("vendor"))
			Expect(hitRecords[0].Header.Get("Authorization")).To(Equal("Bearer sometoken"))

			Expect(searchResults[0].Product.Name).To(Equal("jdk"))
			Expect(searchResults[1].Github.CommitterCount).To(Equal(uint(31337)))
		})
		g.It("should perform a repositories search", func() {
			server.AddPath("/v1/search/repositories").
				SetMethods("GET").
				SetPayload([]byte(sampleValidSourceRepoSearch)).
				SetStatus(http.StatusOK)
			searchResults, err := client.GetSearchRepositories("https://github.com/apache/ode", "sometoken")
			Expect(err).NotTo(HaveOccurred())
			Expect(searchResults).To(HaveLen(1))
			hitRecords := server.HitRecords()
			Expect(hitRecords).To(HaveLen(1))
			Expect(hitRecords[0].Query.Get("github")).To(Equal("https://github.com/apache/ode"))
			Expect(hitRecords[0].Header.Get("Authorization")).To(Equal("Bearer sometoken"))

			Expect(searchResults[0].Product.Name).To(Equal("ode"))
			Expect(searchResults[0].Product.Title).To(Equal("Apache Ode 1.1.1"))
			Expect(searchResults[0].Github.CommitterCount).To(Equal(uint(11)))
		})
	})
}

const (
	sampleValidProdIDSearch     = `{"data": [{"product": {"name": "jdk", "language": "", "title": "IBM Jdk 1.4.2", "created_at": "2017-02-13T20:02:38.33Z", "up": "", "updated_at": "2017-02-13T20:02:38.33Z", "edition": "", "version": "1.4.2", "references": [], "source": null, "org": "ibm", "part": "/a", "external_id": "cpe:/a:ibm:jdk:1.4.2", "id": 52994, "aliases": null}, "github": {"committer_count": 31337, "uri": "https://github.com/ibm/jdk"}}, {"product": {"name": "jdk", "language": "", "title": "IBM Jdk 5.0", "created_at": "2017-02-13T20:02:38.33Z", "up": "", "updated_at": "2017-02-13T20:02:38.33Z", "edition": "", "version": "5.0", "references": [], "source": null, "org": "ibm", "part": "/a", "external_id": "cpe:/a:ibm:jdk:5.0", "id": 52995, "aliases": null}, "github": {"committer_count": 31337, "uri": "https://github.com/ibm/jdk"}}], "meta": { "authors": [ "Ion Channel Dev Team" ], "copyright": "Copyright 2018 Selection Pressure LLC www.selectpress.net", "offset": 0, "total_count": 2, "version": "v1"}}`
	sampleValidSourceRepoSearch = `{"meta": {"total_count": 2, "offset": 0, "version": "v1", "copyright": "Copyright 2018 Selection Pressure LLC www.selectpress.net", "authors": ["Ion Channel Dev Team"]}, "data": [{"product": {"name": "ode", "language": "", "source": null, "created_at": "2018-03-27T22:26:14.991Z", "title": "Apache Ode 1.1.1", "up": "rc1", "updated_at": "2018-03-27T22:26:14.991Z", "edition": "", "part": "/a", "references": null, "version": "1.1.1", "org": "apache", "external_id": "cpe:/a:apache:ode:1.1.1:rc1", "id": 759066869, "aliases": null}, "github": {"committer_count": 11, "uri": "https://github.com/apache/ode"}}]}`
)
<|endoftext|>


<commit_before>package security

import (
	"time"

	"github.com/cuigh/auxo/cache"
	"github.com/cuigh/auxo/log"
	"github.com/cuigh/auxo/net/web"
	"github.com/cuigh/auxo/security"
	"github.com/cuigh/auxo/security/certify"
	"github.com/cuigh/auxo/security/certify/ldap"
	"github.com/cuigh/auxo/security/passwd"
	"github.com/cuigh/swirl/biz"
	"github.com/cuigh/swirl/model"
)

const pkgName = "swirl.security"

// Validator is the authenticator function of swirl.
func Validator(setting *model.Setting) func(name, pwd string) (ticket string, err error) {
	ldapRealm := createLDAPRealm(setting)
	return func(name, pwd string) (ticket string, err error) {
		var (
			su security.User
			mu *model.User
		)

		// try find user first
		mu, err = biz.User.GetByName(name)
		if err != nil {
			return
		}

		if mu != nil && mu.Type == model.UserTypeInternal { // internal user
			if !passwd.Validate(pwd, mu.Password, mu.Salt) {
				err = certify.ErrInvalidToken
			}
		} else if ldapRealm != nil { // user not exist or ldap user
			su, err = ldapRealm.Login(certify.NewSimpleToken(name, pwd))
			if err == nil && mu == nil { // create user if not exists
				lu := su.(*ldap.User)
				mu = &model.User{
					Type:      model.UserTypeLDAP,
					LoginName: lu.LoginName(),
					Name:      lu.Name(),
					Email:     lu.Email(),
				}
				err = biz.User.Create(mu, nil)
			}
		} else {
			err = certify.ErrInvalidToken
		}

		// replace user session, one session is allowed per user.
		if err == nil {
			ticket, err = biz.User.UpdateSession(mu.ID)
		}

		// create event
		if err == nil {
			biz.Event.CreateAuthentication(model.EventActionLogin, mu.ID, mu.LoginName, mu.Name)
		}
		return
	}
}

func createLDAPRealm(setting *model.Setting) certify.Realm {
	if setting.LDAP.Enabled {
		opts := []ldap.Option{
			ldap.NameAttr(setting.LDAP.NameAttr),
			ldap.EmailAttr(setting.LDAP.EmailAttr),
			ldap.UserFilter(setting.LDAP.UserFilter),
			ldap.Security(ldap.SecurityPolicy(setting.LDAP.Security)),
		}
		if setting.LDAP.Authentication == 1 {
			opts = append(opts, ldap.Binding(setting.LDAP.BindDN, setting.LDAP.BindPassword))
		}
		return ldap.New(setting.LDAP.Address, setting.LDAP.BaseDN, setting.LDAP.UserDN, opts...)
	}
	return nil
}

// Identifier is used to identity user.
func Identifier(token string) (user web.User) {
	const cacheKey = "auth_user"

	session, err := biz.User.GetSession(token)
	if err != nil {
		log.Get(pkgName).Errorf("Load session failed: %v", err)
		return
	}
	if session == nil || session.Expires.Before(time.Now()) {
		return
	}

	// try find from cache first
	value := cache.Get(cacheKey, session.UserID)
	if !value.IsNil() {
		user = &model.AuthUser{}
		if err = value.Scan(user); err == nil {
			return
		}
		log.Get(pkgName).Warnf("Load auth user from cache failed: %v", err)
	}

	u, err := biz.User.GetByID(session.UserID)
	if err != nil {
		log.Get(pkgName).Errorf("Load user failed: %v", err)
		return
	}
	if u == nil {
		return
	}

	var roles []*model.Role
	if roles, err = getRoles(u); err == nil {
		user = model.NewAuthUser(u, roles)
		cache.Set(user, cacheKey, session.UserID)
	}
	return
}

func getRoles(u *model.User) (roles []*model.Role, err error) {
	if len(u.Roles) > 0 {
		roles = make([]*model.Role, len(u.Roles))
		for i, id := range u.Roles {
			roles[i], err = biz.Role.Get(id)
			if err != nil {
				return nil, err
			}
		}
	}
	return
}
<commit_msg>Disable login of blocked users<commit_after>package security

import (
	"time"

	"github.com/cuigh/auxo/cache"
	"github.com/cuigh/auxo/log"
	"github.com/cuigh/auxo/net/web"
	"github.com/cuigh/auxo/security"
	"github.com/cuigh/auxo/security/certify"
	"github.com/cuigh/auxo/security/certify/ldap"
	"github.com/cuigh/auxo/security/passwd"
	"github.com/cuigh/swirl/biz"
	"github.com/cuigh/swirl/model"
)

const pkgName = "swirl.security"

// Validator is the authenticator function of swirl.
func Validator(setting *model.Setting) func(name, pwd string) (ticket string, err error) {
	ldapRealm := createLDAPRealm(setting)
	return func(name, pwd string) (ticket string, err error) {
		var (
			su security.User
			mu *model.User
		)

		// try find user first
		mu, err = biz.User.GetByName(name)
		if err != nil {
			return
		}

		if mu != nil && mu.Status == model.UserStatusBlocked {
			return "", certify.ErrAccountDisabled
		}

		if mu != nil && mu.Type == model.UserTypeInternal { // internal user
			if !passwd.Validate(pwd, mu.Password, mu.Salt) {
				err = certify.ErrInvalidToken
			}
		} else if ldapRealm != nil { // user not exist or ldap user
			su, err = ldapRealm.Login(certify.NewSimpleToken(name, pwd))
			if err == nil && mu == nil { // create user if not exists
				lu := su.(*ldap.User)
				mu = &model.User{
					Type:      model.UserTypeLDAP,
					LoginName: lu.LoginName(),
					Name:      lu.Name(),
					Email:     lu.Email(),
				}
				err = biz.User.Create(mu, nil)
			}
		} else {
			err = certify.ErrInvalidToken
		}

		// replace user session, one session is allowed per user.
		if err == nil {
			ticket, err = biz.User.UpdateSession(mu.ID)
		}

		// create event
		if err == nil {
			biz.Event.CreateAuthentication(model.EventActionLogin, mu.ID, mu.LoginName, mu.Name)
		}
		return
	}
}

func createLDAPRealm(setting *model.Setting) certify.Realm {
	if setting.LDAP.Enabled {
		opts := []ldap.Option{
			ldap.NameAttr(setting.LDAP.NameAttr),
			ldap.EmailAttr(setting.LDAP.EmailAttr),
			ldap.UserFilter(setting.LDAP.UserFilter),
			ldap.Security(ldap.SecurityPolicy(setting.LDAP.Security)),
		}
		if setting.LDAP.Authentication == 1 {
			opts = append(opts, ldap.Binding(setting.LDAP.BindDN, setting.LDAP.BindPassword))
		}
		return ldap.New(setting.LDAP.Address, setting.LDAP.BaseDN, setting.LDAP.UserDN, opts...)
	}
	return nil
}

// Identifier is used to identity user.
func Identifier(token string) (user web.User) {
	const cacheKey = "auth_user"

	session, err := biz.User.GetSession(token)
	if err != nil {
		log.Get(pkgName).Errorf("Load session failed: %v", err)
		return
	}
	if session == nil || session.Expires.Before(time.Now()) {
		return
	}

	// try find from cache first
	value := cache.Get(cacheKey, session.UserID)
	if !value.IsNil() {
		user = &model.AuthUser{}
		if err = value.Scan(user); err == nil {
			return
		}
		log.Get(pkgName).Warnf("Load auth user from cache failed: %v", err)
	}

	u, err := biz.User.GetByID(session.UserID)
	if err != nil {
		log.Get(pkgName).Errorf("Load user failed: %v", err)
		return
	}
	if u == nil {
		return
	}

	var roles []*model.Role
	if roles, err = getRoles(u); err == nil {
		user = model.NewAuthUser(u, roles)
		cache.Set(user, cacheKey, session.UserID)
	}
	return
}

func getRoles(u *model.User) (roles []*model.Role, err error) {
	if len(u.Roles) > 0 {
		roles = make([]*model.Role, len(u.Roles))
		for i, id := range u.Roles {
			roles[i], err = biz.Role.Get(id)
			if err != nil {
				return nil, err
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>package htm

import (
	"fmt"
	"github.com/cznic/mathutil"
	"github.com/skelterjohn/go.matrix"
	"math"
	"math/rand"
	"sort"
)
<commit_msg>implemented addToSegmentUpdates()<commit_after>package htm

import (
	"fmt"
	"github.com/cznic/mathutil"
	"github.com/skelterjohn/go.matrix"
	"math"
	"math/rand"
	"sort"
)

type SegmentUpdate struct {
	columnIdx        int
	cellIdx          int
	segment          *Segment
	activeSynapses   []int
	sequenceSegment  bool
	phase1Flag       bool
	weaklyPredicting bool
	segmentUpdates   map[TupleInt][]UpdateState
	lrnIterationIdx  int
}

type UpdateState struct {
	IterationIdx int
	Update       *SegmentUpdate
}

/*
 Store a dated potential segment update. The "date" (iteration index) is used
later to determine whether the update is too old and should be forgotten.
This is controlled by parameter segUpdateValidDuration.
*/

func (su *SegmentUpdate) addToSegmentUpdates(c, i int, segUpdate *SegmentUpdate) {
	if segUpdate == nil || len(segUpdate.activeSynapses) == 0 {
		return
	}

	// key = (column index, cell index in column)
	key := TupleInt{c, i}

	newUpdate := UpdateState{su.lrnIterationIdx, segUpdate}
	if _, ok := su.segmentUpdates[key]; ok {
		su.segmentUpdates[key] = append(su.segmentUpdates[key], newUpdate)
	} else {
		su.segmentUpdates[key] = []UpdateState{newUpdate}
	}

}
<|endoftext|>


<commit_before>package readline

import (
	"fmt"
	"io"
	"os"

	"golang.org/x/crypto/ssh/terminal"
)

type Operation struct {
	cfg     *Config
	t       *Terminal
	buf     *RuneBuffer
	outchan chan []rune
	w       io.Writer

	*opHistory
	*opSearch
	*opCompleter
	*opVim
}

type wrapWriter struct {
	r      *Operation
	t      *Terminal
	target io.Writer
}

func (w *wrapWriter) Write(b []byte) (int, error) {
	if !w.t.IsReading() {
		return w.target.Write(b)
	}

	var (
		n   int
		err error
	)
	w.r.buf.Refresh(func() {
		n, err = w.target.Write(b)
	})

	if w.r.IsSearchMode() {
		w.r.SearchRefresh(-1)
	}
	if w.r.IsInCompleteMode() {
		w.r.CompleteRefresh()
	}
	return n, err
}

func NewOperation(t *Terminal, cfg *Config) *Operation {
	op := &Operation{
		cfg:       cfg,
		t:         t,
		buf:       NewRuneBuffer(t, cfg.Prompt),
		outchan:   make(chan []rune),
		opHistory: newOpHistory(cfg.HistoryFile),
	}
	op.opVim = newVimMode(op)
	op.w = op.buf.w
	op.opSearch = newOpSearch(op.buf.w, op.buf, op.opHistory)
	op.opCompleter = newOpCompleter(op.buf.w, op)
	go op.ioloop()
	return op
}

func (o *Operation) SetPrompt(s string) {
	o.buf.SetPrompt(s)
}

func (o *Operation) ioloop() {
	for {
		keepInSearchMode := false
		keepInCompleteMode := false
		r := o.t.ReadRune()

		if o.IsInCompleteSelectMode() {
			keepInCompleteMode = o.HandleCompleteSelect(r)
			if keepInCompleteMode {
				continue
			}

			o.buf.Refresh(nil)
			switch r {
			case CharEnter, CharCtrlJ:
				o.UpdateHistory(o.buf.Runes(), false)
				fallthrough
			case CharInterrupt:
				o.t.KickRead()
				fallthrough
			case CharBell:
				continue
			}
		}

		if o.IsEnableVimMode() {
			r = o.HandleVim(r, o.t.ReadRune)
			if r == 0 {
				continue
			}
		}

		switch r {
		case CharBell:
			if o.IsSearchMode() {
				o.ExitSearchMode(true)
				o.buf.Refresh(nil)
			}
			if o.IsInCompleteMode() {
				o.ExitCompleteMode(true)
				o.buf.Refresh(nil)
			}
		case CharTab:
			if o.opCompleter == nil {
				break
			}
			o.OnComplete()
			keepInCompleteMode = true
		case CharBckSearch:
			o.SearchMode(S_DIR_BCK)
			keepInSearchMode = true
		case CharCtrlU:
			o.buf.KillFront()
		case CharFwdSearch:
			o.SearchMode(S_DIR_FWD)
			keepInSearchMode = true
		case CharKill:
			o.buf.Kill()
			keepInCompleteMode = true
		case MetaForward:
			o.buf.MoveToNextWord()
		case CharTranspose:
			o.buf.Transpose()
		case MetaBackward:
			o.buf.MoveToPrevWord()
		case MetaDelete:
			o.buf.DeleteWord()
		case CharLineStart:
			o.buf.MoveToLineStart()
		case CharLineEnd:
			o.buf.MoveToLineEnd()
		case CharDelete:
			if !o.buf.Delete() {
				o.t.Bell()
			}
		case CharBackspace, CharCtrlH:
			if o.IsSearchMode() {
				o.SearchBackspace()
				keepInSearchMode = true
				break
			}

			if o.buf.Len() == 0 {
				o.t.Bell()
				break
			}
			o.buf.Backspace()
			if o.IsInCompleteMode() {
				o.OnComplete()
			}
		case MetaBackspace, CharCtrlW:
			o.buf.BackEscapeWord()
		case CharEnter, CharCtrlJ:
			if o.IsSearchMode() {
				o.ExitSearchMode(false)
			}
			o.buf.MoveToLineEnd()
			o.buf.WriteRune('\n')
			data := o.buf.Reset()
			data = data[:len(data)-1] // trim \n
			o.outchan <- data
			o.NewHistory(data)
		case CharBackward:
			o.buf.MoveBackward()
		case CharForward:
			o.buf.MoveForward()
		case CharPrev:
			buf := o.PrevHistory()
			if buf != nil {
				o.buf.Set(buf)
			} else {
				o.t.Bell()
			}
		case CharNext:
			buf, ok := o.NextHistory()
			if ok {
				o.buf.Set(buf)
			} else {
				o.t.Bell()
			}
		case CharInterrupt:
			if o.IsSearchMode() {
				o.t.KickRead()
				o.ExitSearchMode(true)
				break
			}
			if o.IsInCompleteMode() {
				o.t.KickRead()
				o.ExitCompleteMode(true)
				o.buf.Refresh(nil)
				break
			}
			o.buf.MoveToLineEnd()
			o.buf.Refresh(nil)
			o.buf.WriteString("^C\n")
			o.outchan <- nil
		default:
			if o.IsSearchMode() {
				o.SearchChar(r)
				keepInSearchMode = true
				break
			}
			o.buf.WriteRune(r)
			if o.IsInCompleteMode() {
				o.OnComplete()
				keepInCompleteMode = true
			}
		}

		if !keepInSearchMode && o.IsSearchMode() {
			o.ExitSearchMode(false)
			o.buf.Refresh(nil)
		} else if o.IsInCompleteMode() {
			if !keepInCompleteMode {
				o.ExitCompleteMode(false)
				o.buf.Refresh(nil)
			} else {
				o.buf.Refresh(nil)
				o.CompleteRefresh()
			}
		}
		if !o.IsSearchMode() {
			o.UpdateHistory(o.buf.Runes(), false)
		}
	}
}

func (o *Operation) Stderr() io.Writer {
	return &wrapWriter{target: o.cfg.Stderr, r: o, t: o.t}
}

func (o *Operation) Stdout() io.Writer {
	return &wrapWriter{target: o.cfg.Stdout, r: o, t: o.t}
}

func (o *Operation) String() (string, error) {
	r, err := o.Runes()
	if err != nil {
		return "", err
	}
	return string(r), nil
}

func (o *Operation) Runes() ([]rune, error) {
	o.buf.Refresh(nil) // print prompt
	o.t.KickRead()
	r := <-o.outchan
	if r == nil {
		return nil, io.EOF
	}
	return r, nil
}

func (o *Operation) Password(prompt string) ([]byte, error) {
	w := o.Stdout()
	if prompt != "" {
		fmt.Fprintf(w, prompt)
	}
	b, err := terminal.ReadPassword(int(os.Stdin.Fd()))
	fmt.Fprint(w, "\r\n")
	return b, err
}

func (o *Operation) SetTitle(t string) {
	o.w.Write([]byte("\033[2;" + t + "\007"))
}

func (o *Operation) Slice() ([]byte, error) {
	r, err := o.Runes()
	if err != nil {
		return nil, err
	}
	return []byte(string(r)), nil
}

func (o *Operation) Close() {
	o.opHistory.Close()
}
<commit_msg>fixes crash in nil completer #8<commit_after>package readline

import (
	"fmt"
	"io"
	"os"

	"golang.org/x/crypto/ssh/terminal"
)

type Operation struct {
	cfg     *Config
	t       *Terminal
	buf     *RuneBuffer
	outchan chan []rune
	w       io.Writer

	*opHistory
	*opSearch
	*opCompleter
	*opVim
}

type wrapWriter struct {
	r      *Operation
	t      *Terminal
	target io.Writer
}

func (w *wrapWriter) Write(b []byte) (int, error) {
	if !w.t.IsReading() {
		return w.target.Write(b)
	}

	var (
		n   int
		err error
	)
	w.r.buf.Refresh(func() {
		n, err = w.target.Write(b)
	})

	if w.r.IsSearchMode() {
		w.r.SearchRefresh(-1)
	}
	if w.r.IsInCompleteMode() {
		w.r.CompleteRefresh()
	}
	return n, err
}

func NewOperation(t *Terminal, cfg *Config) *Operation {
	op := &Operation{
		cfg:       cfg,
		t:         t,
		buf:       NewRuneBuffer(t, cfg.Prompt),
		outchan:   make(chan []rune),
		opHistory: newOpHistory(cfg.HistoryFile),
	}
	op.opVim = newVimMode(op)
	op.w = op.buf.w
	op.opSearch = newOpSearch(op.buf.w, op.buf, op.opHistory)
	op.opCompleter = newOpCompleter(op.buf.w, op)
	go op.ioloop()
	return op
}

func (o *Operation) SetPrompt(s string) {
	o.buf.SetPrompt(s)
}

func (o *Operation) ioloop() {
	for {
		keepInSearchMode := false
		keepInCompleteMode := false
		r := o.t.ReadRune()

		if o.IsInCompleteSelectMode() {
			keepInCompleteMode = o.HandleCompleteSelect(r)
			if keepInCompleteMode {
				continue
			}

			o.buf.Refresh(nil)
			switch r {
			case CharEnter, CharCtrlJ:
				o.UpdateHistory(o.buf.Runes(), false)
				fallthrough
			case CharInterrupt:
				o.t.KickRead()
				fallthrough
			case CharBell:
				continue
			}
		}

		if o.IsEnableVimMode() {
			r = o.HandleVim(r, o.t.ReadRune)
			if r == 0 {
				continue
			}
		}

		switch r {
		case CharBell:
			if o.IsSearchMode() {
				o.ExitSearchMode(true)
				o.buf.Refresh(nil)
			}
			if o.IsInCompleteMode() {
				o.ExitCompleteMode(true)
				o.buf.Refresh(nil)
			}
		case CharTab:
			if o.cfg.AutoComplete == nil {
				o.t.Bell()
				break
			}
			o.OnComplete()
			keepInCompleteMode = true
		case CharBckSearch:
			o.SearchMode(S_DIR_BCK)
			keepInSearchMode = true
		case CharCtrlU:
			o.buf.KillFront()
		case CharFwdSearch:
			o.SearchMode(S_DIR_FWD)
			keepInSearchMode = true
		case CharKill:
			o.buf.Kill()
			keepInCompleteMode = true
		case MetaForward:
			o.buf.MoveToNextWord()
		case CharTranspose:
			o.buf.Transpose()
		case MetaBackward:
			o.buf.MoveToPrevWord()
		case MetaDelete:
			o.buf.DeleteWord()
		case CharLineStart:
			o.buf.MoveToLineStart()
		case CharLineEnd:
			o.buf.MoveToLineEnd()
		case CharDelete:
			if !o.buf.Delete() {
				o.t.Bell()
			}
		case CharBackspace, CharCtrlH:
			if o.IsSearchMode() {
				o.SearchBackspace()
				keepInSearchMode = true
				break
			}

			if o.buf.Len() == 0 {
				o.t.Bell()
				break
			}
			o.buf.Backspace()
			if o.IsInCompleteMode() {
				o.OnComplete()
			}
		case MetaBackspace, CharCtrlW:
			o.buf.BackEscapeWord()
		case CharEnter, CharCtrlJ:
			if o.IsSearchMode() {
				o.ExitSearchMode(false)
			}
			o.buf.MoveToLineEnd()
			o.buf.WriteRune('\n')
			data := o.buf.Reset()
			data = data[:len(data)-1] // trim \n
			o.outchan <- data
			o.NewHistory(data)
		case CharBackward:
			o.buf.MoveBackward()
		case CharForward:
			o.buf.MoveForward()
		case CharPrev:
			buf := o.PrevHistory()
			if buf != nil {
				o.buf.Set(buf)
			} else {
				o.t.Bell()
			}
		case CharNext:
			buf, ok := o.NextHistory()
			if ok {
				o.buf.Set(buf)
			} else {
				o.t.Bell()
			}
		case CharInterrupt:
			if o.IsSearchMode() {
				o.t.KickRead()
				o.ExitSearchMode(true)
				break
			}
			if o.IsInCompleteMode() {
				o.t.KickRead()
				o.ExitCompleteMode(true)
				o.buf.Refresh(nil)
				break
			}
			o.buf.MoveToLineEnd()
			o.buf.Refresh(nil)
			o.buf.WriteString("^C\n")
			o.outchan <- nil
		default:
			if o.IsSearchMode() {
				o.SearchChar(r)
				keepInSearchMode = true
				break
			}
			o.buf.WriteRune(r)
			if o.IsInCompleteMode() {
				o.OnComplete()
				keepInCompleteMode = true
			}
		}

		if !keepInSearchMode && o.IsSearchMode() {
			o.ExitSearchMode(false)
			o.buf.Refresh(nil)
		} else if o.IsInCompleteMode() {
			if !keepInCompleteMode {
				o.ExitCompleteMode(false)
				o.buf.Refresh(nil)
			} else {
				o.buf.Refresh(nil)
				o.CompleteRefresh()
			}
		}
		if !o.IsSearchMode() {
			o.UpdateHistory(o.buf.Runes(), false)
		}
	}
}

func (o *Operation) Stderr() io.Writer {
	return &wrapWriter{target: o.cfg.Stderr, r: o, t: o.t}
}

func (o *Operation) Stdout() io.Writer {
	return &wrapWriter{target: o.cfg.Stdout, r: o, t: o.t}
}

func (o *Operation) String() (string, error) {
	r, err := o.Runes()
	if err != nil {
		return "", err
	}
	return string(r), nil
}

func (o *Operation) Runes() ([]rune, error) {
	o.buf.Refresh(nil) // print prompt
	o.t.KickRead()
	r := <-o.outchan
	if r == nil {
		return nil, io.EOF
	}
	return r, nil
}

func (o *Operation) Password(prompt string) ([]byte, error) {
	w := o.Stdout()
	if prompt != "" {
		fmt.Fprintf(w, prompt)
	}
	b, err := terminal.ReadPassword(int(os.Stdin.Fd()))
	fmt.Fprint(w, "\r\n")
	return b, err
}

func (o *Operation) SetTitle(t string) {
	o.w.Write([]byte("\033[2;" + t + "\007"))
}

func (o *Operation) Slice() ([]byte, error) {
	r, err := o.Runes()
	if err != nil {
		return nil, err
	}
	return []byte(string(r)), nil
}

func (o *Operation) Close() {
	o.opHistory.Close()
}
<|endoftext|>