first commit

2026-08-12 21:35:40 +08:00 · 2016-03-17 20:17:59 -04:00
commit 2f232c194a
27 changed files with 1816 additions and 0 deletions
--- a/command/datastore.go
+++ b/command/datastore.go
@@ -0,0 +1,173 @@
+package command
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/garyburd/redigo/redis"
+	"github.com/go-kit/kit/log"
+	"github.com/groob/plist"
+	"github.com/micromdm/mdm"
+)
+
+var (
+	// ErrNoKey is returned if there is no key in redis
+	ErrNoKey = errors.New("There is no such key in redis.")
+)
+
+// Datastore manages MDM Payloads in redis
+type Datastore interface {
+	// Saves the payload in redis
+	// SET CommandUUID plistData
+	SavePayload(payload *mdm.Payload) error
+	// Adds MDM commands to a queue in redis list
+	// LPUSH deviceUDID commandUUID
+	QueueCommand(deviceUDID, commandUUID string) error
+	NextCommand(deviceUDID string) ([]byte, int, error)
+	DeleteCommand(deviceUDID, commandUUID string) (int, error)
+}
+
+type redisDB struct {
+	pool *redis.Pool
+}
+
+// NewDB creates a new databases connection
+func NewDB(driver, conn string, options ...func(*config) error) Datastore {
+	conf := &config{}
+	defaultLogger := log.NewLogfmtLogger(os.Stderr)
+	for _, option := range options {
+		if err := option(conf); err != nil {
+			defaultLogger.Log("err", err)
+			os.Exit(1)
+		}
+	}
+	switch driver {
+	case "redis":
+		return redisDB{pool: redisPool(conn, conf.logger)}
+	default:
+		conf.logger.Log("err", "unknown driver")
+		os.Exit(1)
+		return nil
+	}
+}
+
+func redisPool(conn string, logger log.Logger) *redis.Pool {
+	pool := &redis.Pool{
+		MaxIdle:     3,
+		IdleTimeout: 240 * time.Second,
+		Dial: func() (redis.Conn, error) {
+			c, err := redis.Dial("tcp", conn)
+			if err != nil {
+				return nil, err
+			}
+			return c, err
+		},
+		TestOnBorrow: func(c redis.Conn, t time.Time) error {
+			_, err := c.Do("PING")
+			return err
+		},
+	}
+	checkRedisConn(pool, logger)
+	return pool
+}
+
+func checkRedisConn(pool *redis.Pool, logger log.Logger) {
+	conn := pool.Get()
+	defer conn.Close()
+
+	var dbError error
+	maxAttempts := 20
+	for attempts := 1; attempts <= maxAttempts; attempts++ {
+		_, dbError = conn.Do("PING")
+		if dbError == nil {
+			break
+		}
+		logger.Log("msg", fmt.Sprintf("could not connect to redis: %v", dbError))
+		time.Sleep(time.Duration(attempts) * time.Second)
+	}
+	if dbError != nil {
+		logger.Log("err", dbError)
+		os.Exit(1)
+	}
+}
+
+func (rds redisDB) SavePayload(payload *mdm.Payload) error {
+	var buf bytes.Buffer
+	// get connection from redis pool
+	conn := rds.pool.Get()
+	defer conn.Close()
+	// encode payload into a plist
+	err := plist.NewEncoder(&buf).Encode(payload)
+	if err != nil {
+		return err
+	}
+	// create a commandUUID key with the plist as the value
+	_, err = conn.Do("set", payload.CommandUUID, buf.String())
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (rds redisDB) QueueCommand(deviceUDID, commandUUID string) error {
+	// get connection from redis pool
+	conn := rds.pool.Get()
+	defer conn.Close()
+	_, err := conn.Do("lpush", deviceUDID, commandUUID)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+func (rds redisDB) NextCommand(deviceUDID string) ([]byte, int, error) {
+	// get connection from redis pool
+	conn := rds.pool.Get()
+	defer conn.Close()
+	// pop the first command
+	commandUUID, err := redis.String(conn.Do("lpop", deviceUDID))
+	if err != nil && err != redis.ErrNil {
+		return nil, 0, err
+	}
+	// if the list is empty
+	if err == redis.ErrNil {
+		return []byte{}, 0, nil
+	}
+	// push the redis command back to the end of the list
+	_, err = conn.Do("rpush", deviceUDID, commandUUID)
+	command, err := redis.String(conn.Do("get", commandUUID))
+	if err == redis.ErrNil {
+		return nil, 0, ErrNoKey
+	}
+
+	// get a command list length
+	total, err := redis.Int(conn.Do("llen", deviceUDID))
+	if err != nil {
+		return nil, 0, err
+	}
+	return []byte(command), total, err
+}
+
+func (rds redisDB) DeleteCommand(deviceUDID, commandUUID string) (int, error) {
+	// get connection from redis pool
+	conn := rds.pool.Get()
+	defer conn.Close()
+	// remove from list
+	_, err := conn.Do("lrem", deviceUDID, 0, commandUUID)
+	if err != nil {
+		return 0, err
+	}
+	// set the key to expire in an hour
+	_, err = conn.Do("expire", commandUUID, 3600)
+	if err != nil {
+		return 0, err
+	}
+	// get a command list length
+	total, err := redis.Int(conn.Do("llen", deviceUDID))
+	if err != nil {
+		return 0, err
+	}
+	return total, nil
+}
--- a/command/service.go
+++ b/command/service.go
@@ -0,0 +1,128 @@
+package command
+
+import (
+	"net/http"
+	"os"
+
+	"golang.org/x/net/context"
+
+	httptransport "github.com/go-kit/kit/transport/http"
+
+	"github.com/go-kit/kit/log"
+	"github.com/gorilla/mux"
+	"github.com/micromdm/mdm"
+)
+
+// MDMCommandService allows creating and deleting MDM Command Payloads
+type MDMCommandService interface {
+	NewCommand(*mdm.CommandRequest) (*mdm.Payload, error)
+	NextCommand(udid string) ([]byte, int, error)
+	DeleteCommand(deviceUDID, commandUUID string) (int, error)
+}
+
+type mdmCommandService struct {
+	// a redis datastore
+	db Datastore
+}
+
+type config struct {
+	logger log.Logger
+	db     Datastore
+}
+
+// NewCommandService creates a new MDM Command Service
+func NewCommandService(options ...func(*config) error) MDMCommandService {
+	conf := &config{}
+	defaultLogger := log.NewLogfmtLogger(os.Stderr)
+	for _, option := range options {
+		if err := option(conf); err != nil {
+			defaultLogger.Log("err", err)
+			os.Exit(1)
+		}
+	}
+	var svc MDMCommandService
+	svc = mdmCommandService{db: conf.db}
+	return svc
+}
+
+// Logger adds a logger to the service
+func Logger(logger log.Logger) func(*config) error {
+	return func(c *config) error {
+		c.logger = logger
+		return nil
+	}
+}
+
+// DB adds a db connection to the service
+func DB(db Datastore) func(*config) error {
+	return func(c *config) error {
+		c.db = db
+		return nil
+	}
+}
+
+// ServiceHandler returns an http handler for the command service
+func ServiceHandler(ctx context.Context, svc MDMCommandService) http.Handler {
+	commonOptions := []httptransport.ServerOption{
+		httptransport.ServerErrorEncoder(encodeError),
+	}
+	newCommandEndpoint := makeNewCommandEndpoint(svc)
+	newCommandHandler := httptransport.NewServer(
+		ctx,
+		newCommandEndpoint,
+		decodeNewCommandRequest,
+		encodeResponse,
+		commonOptions...,
+	)
+	nextCommandEndpoint := makeNextCommandEndpoint(svc)
+	nextCommandHandler := httptransport.NewServer(
+		ctx,
+		nextCommandEndpoint,
+		decodeNextCommandRequest,
+		encodeResponse,
+		commonOptions...,
+	)
+	deleteCommandEndpoint := makeDeleteCommandEndpoint(svc)
+	deleteCommandHandler := httptransport.NewServer(
+		ctx,
+		deleteCommandEndpoint,
+		decodeDeleteCommandRequest,
+		encodeResponse,
+		commonOptions...,
+	)
+	r := mux.NewRouter()
+	r.Methods("POST").Path("/mdm/commands").Handler(newCommandHandler)
+	r.Methods("GET").Path("/mdm/commands/{udid}/next").Handler(nextCommandHandler)
+	r.Methods("DELETE").Path("/mdm/commands/{udid}/{uuid}").Handler(deleteCommandHandler)
+	return r
+}
+
+func (svc mdmCommandService) NewCommand(request *mdm.CommandRequest) (*mdm.Payload, error) {
+	// create a payload
+	payload, err := mdm.NewPayload(request)
+	if err != nil {
+		return nil, err
+	}
+	// save in redis
+	err = svc.db.SavePayload(payload)
+	if err != nil {
+		return nil, err
+	}
+	// add command to a queue in redis
+	err = svc.db.QueueCommand(request.UDID, payload.CommandUUID)
+	if err != nil {
+		return nil, err
+	}
+	// return created payload to user
+	return payload, nil
+}
+
+// NextCommand returns an MDM Payload from a list of queued payloads
+func (svc mdmCommandService) NextCommand(udid string) ([]byte, int, error) {
+	return svc.db.NextCommand(udid)
+}
+
+// DeleteCommand returns an MDM Payload from a list of queued payloads
+func (svc mdmCommandService) DeleteCommand(deviceUDID, commandUUID string) (int, error) {
+	return svc.db.DeleteCommand(deviceUDID, commandUUID)
+}
--- a/command/transport.go
+++ b/command/transport.go
@@ -0,0 +1,251 @@
+package command
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"log"
+	"net/http"
+
+	"github.com/go-kit/kit/endpoint"
+	"github.com/gorilla/mux"
+	"github.com/micromdm/mdm"
+	"golang.org/x/net/context"
+)
+
+var (
+	// ErrEmptyRequest is returned if the request body is empty
+	ErrEmptyRequest = errors.New("Request must contain UDID of the device")
+	errBadRouting   = errors.New("inconsistent mapping between route and handler (programmer error)")
+)
+
+// NewCommandRequest represents an HTTP Request for a new MDM Command
+type NewCommandRequest struct {
+	*mdm.CommandRequest
+}
+
+func decodeNewCommandRequest(r *http.Request) (interface{}, error) {
+	var request NewCommandRequest
+	err := json.NewDecoder(r.Body).Decode(&request.CommandRequest)
+	return request, err
+}
+
+// NewCommandResponse is a command reponse
+type NewCommandResponse struct {
+	*mdm.Payload
+	Err error `json:"error,omitempty"`
+}
+
+func (r NewCommandResponse) error() error { return r.Err }
+
+// errorer is implemented by all concrete response types. It allows us to
+// change the HTTP response code without needing to trigger an endpoint
+// (transport-level) error. For more information, read the big comment in
+// endpoint.go.
+type errorer interface {
+	error() error
+}
+
+// NextCommandRequest is a request to return the next command in a device queue
+type NextCommandRequest struct {
+	UDID string
+}
+
+func decodeNextCommandRequest(r *http.Request) (interface{}, error) {
+	vars := mux.Vars(r)
+	udid, ok := vars["udid"]
+	if !ok {
+		return nil, errBadRouting
+	}
+	var request NextCommandRequest
+	request.UDID = udid
+	return request, nil
+}
+
+// EncodeNextCommandRequest encodes a request for the NextCommand endpoint
+func EncodeNextCommandRequest(r *http.Request, request interface{}) error {
+	req := request.(NextCommandRequest)
+	path := r.URL.Path
+	r.URL.Path = fmt.Sprintf("%v/%v/next", path, req.UDID)
+	var buf bytes.Buffer
+	if err := json.NewEncoder(&buf).Encode(request); err != nil {
+		return err
+	}
+	r.Body = ioutil.NopCloser(&buf)
+	return nil
+}
+
+// NextCommandResponse is a response for the next command
+type NextCommandResponse struct {
+	Payload []byte `json:"command_payload"`
+	Total   int    `json:"total_payloads"`
+	Err     error  `json:"error,omitempty"`
+}
+
+// DecodeNextCommandResponse decodes the response from the provided HTTP response,
+// simply by JSON decoding from the response body. It's designed to be used in
+// transport/http.Client.
+// first decode into map[string]interface{} and check for error in the response
+func DecodeNextCommandResponse(resp *http.Response) (interface{}, error) {
+	var r map[string]interface{}
+	var response NextCommandResponse
+	err := json.NewDecoder(resp.Body).Decode(&r)
+	if rs, ok := r["error"]; ok {
+		response.Err = errors.New(rs.(string))
+	}
+	if rs, ok := r["total_payloads"]; ok {
+		response.Total = int(rs.(float64))
+	}
+	if rs, ok := r["command_payload"]; ok {
+		response.Payload = []byte(rs.(string))
+	}
+	return response, err
+}
+
+func (r NextCommandResponse) error() error { return r.Err }
+
+// DeleteCommandRequest is a request to delete a command
+type DeleteCommandRequest struct {
+	// device UDID
+	UDID string
+	// command UUID
+	UUID string
+}
+
+// EncodeDeleteCommandRequest encodes a request for the NextCommand endpoint
+func EncodeDeleteCommandRequest(r *http.Request, request interface{}) error {
+	req := request.(DeleteCommandRequest)
+	path := r.URL.Path
+	r.URL.Path = fmt.Sprintf("%v/%v/%v", path, req.UDID, req.UUID)
+	var buf bytes.Buffer
+	if err := json.NewEncoder(&buf).Encode(request); err != nil {
+		return err
+	}
+	r.Body = ioutil.NopCloser(&buf)
+	return nil
+}
+
+func decodeDeleteCommandRequest(r *http.Request) (interface{}, error) {
+	vars := mux.Vars(r)
+	udid, ok := vars["udid"]
+	if !ok {
+		return nil, errBadRouting
+	}
+	uuid, ok := vars["uuid"]
+	if !ok {
+		return nil, errBadRouting
+	}
+	var request DeleteCommandRequest
+	request.UDID = udid
+	request.UUID = uuid
+	return request, nil
+}
+
+// DeleteCommandResponse is a response for a delete request
+type DeleteCommandResponse struct {
+	Total int   `json:"remaining_payloads"`
+	Err   error `json:"error,omitempty"`
+}
+
+// DecodeDeleteCommandResponse decodes the response from the provided HTTP response,
+// simply by JSON decoding from the response body. It's designed to be used in
+// transport/http.Client.
+// first decode into map[string]interface{} and check for error in the response
+func DecodeDeleteCommandResponse(resp *http.Response) (interface{}, error) {
+	var r map[string]interface{}
+	var response DeleteCommandResponse
+	err := json.NewDecoder(resp.Body).Decode(&r)
+	if rs, ok := r["error"]; ok {
+		response.Err = errors.New(rs.(string))
+	}
+	if rs, ok := r["remaining_payloads"]; ok {
+		response.Total = int(rs.(float64))
+	}
+	return response, err
+}
+
+func makeNewCommandEndpoint(svc MDMCommandService) endpoint.Endpoint {
+	return func(ctx context.Context, request interface{}) (interface{}, error) {
+		req := request.(NewCommandRequest)
+		if req.UDID == "" || req.RequestType == "" {
+			return NewCommandResponse{Err: ErrEmptyRequest}, nil
+		}
+		payload, err := svc.NewCommand(req.CommandRequest)
+		if err != nil {
+			return NewCommandResponse{Err: err}, nil
+		}
+		return NewCommandResponse{Payload: payload}, nil
+	}
+}
+
+func makeNextCommandEndpoint(svc MDMCommandService) endpoint.Endpoint {
+	return func(ctx context.Context, request interface{}) (interface{}, error) {
+		req := request.(NextCommandRequest)
+		if req.UDID == "" {
+			return NextCommandResponse{Err: ErrEmptyRequest}, nil
+		}
+		payload, total, err := svc.NextCommand(req.UDID)
+		if err != nil {
+			return NextCommandResponse{Err: err}, nil
+		}
+		return NextCommandResponse{Payload: payload, Total: total}, nil
+	}
+}
+
+func makeDeleteCommandEndpoint(svc MDMCommandService) endpoint.Endpoint {
+	return func(ctx context.Context, request interface{}) (interface{}, error) {
+		req := request.(DeleteCommandRequest)
+		if req.UDID == "" {
+			return DeleteCommandResponse{Err: ErrEmptyRequest}, nil
+		}
+		if req.UUID == "" {
+			return DeleteCommandResponse{Err: ErrEmptyRequest}, nil
+		}
+		total, err := svc.DeleteCommand(req.UDID, req.UUID)
+		if err != nil {
+			return DeleteCommandResponse{Err: err}, nil
+		}
+		return DeleteCommandResponse{Total: total}, nil
+	}
+}
+
+// encodeResponse is the common method to encode all response types to the
+// client. I chose to do it this way because I didn't know if something more
+// specific was necessary. It's certainly possible to specialize on a
+// per-response (per-method) basis.
+func encodeResponse(w http.ResponseWriter, response interface{}) error {
+	if e, ok := response.(errorer); ok && e.error() != nil {
+		// Not a Go kit transport error, but a business-logic error.
+		// Provide those as HTTP errors.
+		encodeError(w, e.error())
+		return nil
+	}
+	jsn, err := json.MarshalIndent(response, "", "  ")
+	if err != nil {
+		return err
+	}
+	w.Write(jsn)
+	return nil
+}
+
+func encodeError(w http.ResponseWriter, err error) {
+	w.WriteHeader(codeFrom(err))
+	response := map[string]interface{}{
+		"error": err.Error(),
+	}
+	jsn, err := json.MarshalIndent(response, "", "  ")
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	w.Write(jsn)
+}
+
+func codeFrom(err error) int {
+	switch err {
+	default:
+		return http.StatusInternalServerError
+	}
+}