package command import ( "bytes" "errors" "fmt" "os" "time" "github.com/garyburd/redigo/redis" kitlog "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 provides methods for saving and retrieving MDM commands 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) Commands(deviceUDID string) ([]mdm.Payload, error) Find(commandUUID string) (*mdm.Payload, error) } //NewDB creates a Datastore func NewDB(driver, conn string, logger kitlog.Logger) (Datastore, error) { var ds Datastore switch driver { case "redis": ds = redisDB{pool: redisPool(conn, logger)} return ds, nil default: return nil, errors.New("unknown driver") } } type redisDB struct { pool *redis.Pool } 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 } func (rds redisDB) Commands(deviceUDID string) ([]mdm.Payload, error) { conn := rds.pool.Get() defer conn.Close() commandUUIDs, err := redis.Values(conn.Do("LRANGE", deviceUDID, "0", "-1")) if err != nil { return nil, err } var payloads []mdm.Payload = make([]mdm.Payload, len(commandUUIDs)) for i, commandUUID := range commandUUIDs { payloadData, err := redis.Bytes(conn.Do("GET", commandUUID)) if err != nil { return nil, err } if err := plist.NewDecoder(bytes.NewReader(payloadData)).Decode(&payloads[i]); err != nil { return nil, err } } return payloads, nil } func (rds redisDB) Find(commandUUID string) (*mdm.Payload, error) { conn := rds.pool.Get() defer conn.Close() payloadData, err := redis.Bytes(conn.Do("GET", commandUUID)) if err != nil { return nil, err } var payload *mdm.Payload if err := plist.NewDecoder(bytes.NewReader(payloadData)).Decode(&payload); err != nil { return nil, err } return payload, nil } func redisPool(conn string, logger kitlog.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 kitlog.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) } }