adapter.go 8.98 KB
Newer Older
1 2 3
package pgadapter

import (
khoipham's avatar
khoipham committed
4
	"fmt"
5
	"strings"
khoipham's avatar
khoipham committed
6 7 8 9 10

	"github.com/casbin/casbin/v2/model"
	"github.com/casbin/casbin/v2/persist"
	"github.com/go-pg/pg/v9"
	"github.com/go-pg/pg/v9/orm"
11
	"github.com/go-pg/pg/v9/types"
12
	"github.com/mmcloughlin/meow"
13 14 15 16
)

// CasbinRule represents a rule in Casbin.
type CasbinRule struct {
khoipham's avatar
khoipham committed
17
	ID    string
18 19 20 21 22 23 24 25 26
	PType string
	V0    string
	V1    string
	V2    string
	V3    string
	V4    string
	V5    string
}

27 28 29 30 31
type Filter struct {
	P []string
	G []string
}

32 33
// Adapter represents the github.com/go-pg/pg adapter for policy storage.
type Adapter struct {
34 35 36
	db        *pg.DB
	tableName string
	filtered  bool
37 38
}

39 40
type Option func(a *Adapter)

41
// NewAdapter is the constructor for Adapter.
khoipham's avatar
khoipham committed
42 43 44 45
// arg should be a PostgreS URL string or of type *pg.Options
// The adapter will create a DB named "casbin" if it doesn't exist
func NewAdapter(arg interface{}) (*Adapter, error) {
	db, err := createCasbinDatabase(arg)
46
	if err != nil {
khoipham's avatar
khoipham committed
47
		return nil, fmt.Errorf("pgadapter.NewAdapter: %v", err)
48 49
	}

khoipham's avatar
khoipham committed
50
	a := &Adapter{db: db}
51

52
	if err := a.createTableifNotExists(); err != nil {
khoipham's avatar
khoipham committed
53
		return nil, fmt.Errorf("pgadapter.NewAdapter: %v", err)
54 55
	}

khoipham's avatar
khoipham committed
56
	return a, nil
57 58
}

59 60
// NewAdapterByDB creates new Adapter by using existing DB connection
// creates table from CasbinRule struct if it doesn't exist
61
func NewAdapterByDB(db *pg.DB, opts ...Option) (*Adapter, error) {
62
	a := &Adapter{db: db}
63 64 65 66 67 68 69 70 71 72
	for _, opt := range opts {
		opt(a)
	}

	if len(a.tableName) > 0 {
		a.db.Model((*CasbinRule)(nil)).TableModel().Table().Name = a.tableName
		a.db.Model((*CasbinRule)(nil)).TableModel().Table().FullName = (types.Safe)(a.tableName)
		a.db.Model((*CasbinRule)(nil)).TableModel().Table().FullNameForSelects = (types.Safe)(a.tableName)
	}

73
	if err := a.createTableifNotExists(); err != nil {
74 75 76 77 78
		return nil, fmt.Errorf("pgadapter.NewAdapter: %v", err)
	}
	return a, nil
}

79 80 81 82 83 84 85
// WithTableName can be used to pass custom table name for Casbin rules
func WithTableName(tableName string) Option {
	return func(a *Adapter) {
		a.tableName = tableName
	}
}

khoipham's avatar
khoipham committed
86 87
func createCasbinDatabase(arg interface{}) (*pg.DB, error) {
	var opts *pg.Options
88
	var err error
khoipham's avatar
khoipham committed
89 90 91 92 93 94 95 96 97 98
	if connURL, ok := arg.(string); ok {
		opts, err = pg.ParseURL(connURL)
		if err != nil {
			return nil, err
		}
	} else {
		opts, ok = arg.(*pg.Options)
		if !ok {
			return nil, fmt.Errorf("must pass in a PostgreS URL string or an instance of *pg.Options, received %T instead", arg)
		}
99 100
	}

khoipham's avatar
khoipham committed
101 102
	db := pg.Connect(opts)
	defer db.Close()
103

khoipham's avatar
khoipham committed
104 105 106 107 108
	_, err = db.Exec("CREATE DATABASE casbin")
	db.Close()

	opts.Database = "casbin"
	db = pg.Connect(opts)
109

khoipham's avatar
khoipham committed
110
	return db, nil
111 112
}

khoipham's avatar
khoipham committed
113 114 115 116
// Close close database connection
func (a *Adapter) Close() error {
	if a != nil && a.db != nil {
		return a.db.Close()
117 118 119 120
	}
	return nil
}

121
func (a *Adapter) createTableifNotExists() error {
122
	err := a.db.CreateTable(&CasbinRule{}, &orm.CreateTableOptions{
123 124
		Temp:        false,
		IfNotExists: true,
125 126
	})
	if err != nil {
127
		return err
128 129 130 131
	}
	return nil
}

132
func (r *CasbinRule) String() string {
133 134 135
	const prefixLine = ", "
	var sb strings.Builder

136 137
	sb.WriteString(r.PType)
	if len(r.V0) > 0 {
138
		sb.WriteString(prefixLine)
139
		sb.WriteString(r.V0)
140
	}
141
	if len(r.V1) > 0 {
142
		sb.WriteString(prefixLine)
143
		sb.WriteString(r.V1)
144
	}
145
	if len(r.V2) > 0 {
146
		sb.WriteString(prefixLine)
147
		sb.WriteString(r.V2)
148
	}
149
	if len(r.V3) > 0 {
150
		sb.WriteString(prefixLine)
151
		sb.WriteString(r.V3)
152
	}
153
	if len(r.V4) > 0 {
154
		sb.WriteString(prefixLine)
155
		sb.WriteString(r.V4)
156
	}
157
	if len(r.V5) > 0 {
158
		sb.WriteString(prefixLine)
159
		sb.WriteString(r.V5)
160 161
	}

162
	return sb.String()
163 164 165 166 167 168
}

// LoadPolicy loads policy from database.
func (a *Adapter) LoadPolicy(model model.Model) error {
	var lines []*CasbinRule

169
	if err := a.db.Model(&lines).Select(); err != nil {
170 171 172 173
		return err
	}

	for _, line := range lines {
174
		persist.LoadPolicyLine(line.String(), model)
175 176
	}

177 178
	a.filtered = false

179 180 181
	return nil
}

khoipham's avatar
khoipham committed
182 183
func policyID(ptype string, rule []string) string {
	data := strings.Join(append([]string{ptype}, rule...), ",")
184
	sum := meow.Checksum(0, []byte(data))
khoipham's avatar
khoipham committed
185 186 187
	return fmt.Sprintf("%x", sum)
}

188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210
func savePolicyLine(ptype string, rule []string) *CasbinRule {
	line := &CasbinRule{PType: ptype}

	l := len(rule)
	if l > 0 {
		line.V0 = rule[0]
	}
	if l > 1 {
		line.V1 = rule[1]
	}
	if l > 2 {
		line.V2 = rule[2]
	}
	if l > 3 {
		line.V3 = rule[3]
	}
	if l > 4 {
		line.V4 = rule[4]
	}
	if l > 5 {
		line.V5 = rule[5]
	}

khoipham's avatar
khoipham committed
211
	line.ID = policyID(ptype, rule)
khoipham's avatar
khoipham committed
212

213 214 215 216 217
	return line
}

// SavePolicy saves policy to database.
func (a *Adapter) SavePolicy(model model.Model) error {
218 219 220 221 222
	_, err := a.db.Model((*CasbinRule)(nil)).Where("id IS NOT NULL").Delete()
	if err != nil {
		return err
	}

223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238
	var lines []*CasbinRule

	for ptype, ast := range model["p"] {
		for _, rule := range ast.Policy {
			line := savePolicyLine(ptype, rule)
			lines = append(lines, line)
		}
	}

	for ptype, ast := range model["g"] {
		for _, rule := range ast.Policy {
			line := savePolicyLine(ptype, rule)
			lines = append(lines, line)
		}
	}

239 240 241 242 243
	if len(lines) > 0 {
		_, err = a.db.Model(&lines).
			OnConflict("DO NOTHING").
			Insert()
		return err
DivyPatel9881's avatar
DivyPatel9881 committed
244
	}
245
	return nil
246 247 248 249 250
}

// AddPolicy adds a policy rule to the storage.
func (a *Adapter) AddPolicy(sec string, ptype string, rule []string) error {
	line := savePolicyLine(ptype, rule)
khoipham's avatar
khoipham committed
251 252 253
	_, err := a.db.Model(line).
		OnConflict("DO NOTHING").
		Insert()
254 255 256
	return err
}

257 258 259
// AddPolicies adds policy rules to the storage.
func (a *Adapter) AddPolicies(sec string, ptype string, rules [][]string) error {
	var lines []*CasbinRule
anton troyanov's avatar
anton troyanov committed
260
	for _, rule := range rules {
261 262 263
		line := savePolicyLine(ptype, rule)
		lines = append(lines, line)
	}
anton troyanov's avatar
anton troyanov committed
264

265
	err := a.db.RunInTransaction(func(tx *pg.Tx) error {
DivyPatel9881's avatar
DivyPatel9881 committed
266
		_, err := tx.Model(&lines).
anton troyanov's avatar
anton troyanov committed
267 268
			OnConflict("DO NOTHING").
			Insert()
DivyPatel9881's avatar
DivyPatel9881 committed
269
		return err
270
	})
anton troyanov's avatar
anton troyanov committed
271

272 273 274
	return err
}

275 276 277 278 279 280 281
// RemovePolicy removes a policy rule from the storage.
func (a *Adapter) RemovePolicy(sec string, ptype string, rule []string) error {
	line := savePolicyLine(ptype, rule)
	err := a.db.Delete(line)
	return err
}

282 283 284
// RemovePolicies removes policy rules from the storage.
func (a *Adapter) RemovePolicies(sec string, ptype string, rules [][]string) error {
	var lines []*CasbinRule
anton troyanov's avatar
anton troyanov committed
285
	for _, rule := range rules {
286 287 288 289 290
		line := savePolicyLine(ptype, rule)
		lines = append(lines, line)
	}

	err := a.db.RunInTransaction(func(tx *pg.Tx) error {
DivyPatel9881's avatar
DivyPatel9881 committed
291
		_, err := tx.Model(&lines).
anton troyanov's avatar
anton troyanov committed
292
			Delete()
DivyPatel9881's avatar
DivyPatel9881 committed
293
		return err
294
	})
anton troyanov's avatar
anton troyanov committed
295

296 297 298
	return err
}

299 300
// RemoveFilteredPolicy removes policy rules that match the filter from the storage.
func (a *Adapter) RemoveFilteredPolicy(sec string, ptype string, fieldIndex int, fieldValues ...string) error {
khoipham's avatar
khoipham committed
301
	query := a.db.Model((*CasbinRule)(nil)).Where("p_type = ?", ptype)
302 303

	idx := fieldIndex + len(fieldValues)
khoipham's avatar
khoipham committed
304 305
	if fieldIndex <= 0 && idx > 0 && fieldValues[0-fieldIndex] != "" {
		query = query.Where("v0 = ?", fieldValues[0-fieldIndex])
306
	}
khoipham's avatar
khoipham committed
307 308
	if fieldIndex <= 1 && idx > 1 && fieldValues[1-fieldIndex] != "" {
		query = query.Where("v1 = ?", fieldValues[1-fieldIndex])
309
	}
khoipham's avatar
khoipham committed
310 311
	if fieldIndex <= 2 && idx > 2 && fieldValues[2-fieldIndex] != "" {
		query = query.Where("v2 = ?", fieldValues[2-fieldIndex])
312
	}
khoipham's avatar
khoipham committed
313 314
	if fieldIndex <= 3 && idx > 3 && fieldValues[3-fieldIndex] != "" {
		query = query.Where("v3 = ?", fieldValues[3-fieldIndex])
315
	}
khoipham's avatar
khoipham committed
316 317
	if fieldIndex <= 4 && idx > 4 && fieldValues[4-fieldIndex] != "" {
		query = query.Where("v4 = ?", fieldValues[4-fieldIndex])
318
	}
khoipham's avatar
khoipham committed
319 320
	if fieldIndex <= 5 && idx > 5 && fieldValues[5-fieldIndex] != "" {
		query = query.Where("v5 = ?", fieldValues[5-fieldIndex])
321 322
	}

khoipham's avatar
khoipham committed
323
	_, err := query.Delete()
324 325
	return err
}
326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409

func (a *Adapter) LoadFilteredPolicy(model model.Model, filter interface{}) error {
	if filter == nil {
		return a.LoadPolicy(model)
	}

	filterValue, ok := filter.(*Filter)
	if !ok {
		return fmt.Errorf("invalid filter type")
	}
	err := a.loadFilteredPolicy(model, filterValue, persist.LoadPolicyLine)
	if err != nil {
		return err
	}
	a.filtered = true
	return nil
}

func buildQuery(query *orm.Query, values []string) (*orm.Query, error) {
	for ind, v := range values {
		if v == "" {
			continue
		}
		switch ind {
		case 0:
			query = query.Where("v0 = ?", v)
		case 1:
			query = query.Where("v1 = ?", v)
		case 2:
			query = query.Where("v2 = ?", v)
		case 3:
			query = query.Where("v3 = ?", v)
		case 4:
			query = query.Where("v4 = ?", v)
		case 5:
			query = query.Where("v5 = ?", v)
		default:
			return nil, fmt.Errorf("filter has more values than expected, should not exceed 6 values")
		}
	}
	return query, nil
}

func (a *Adapter) loadFilteredPolicy(model model.Model, filter *Filter, handler func(string, model.Model)) error {
	if filter.P != nil {
		lines := []*CasbinRule{}

		query := a.db.Model(&lines).Where("p_type = 'p'")
		query, err := buildQuery(query, filter.P)
		if err != nil {
			return err
		}
		err = query.Select()
		if err != nil {
			return err
		}

		for _, line := range lines {
			handler(line.String(), model)
		}
	}
	if filter.G != nil {
		lines := []*CasbinRule{}

		query := a.db.Model(&lines).Where("p_type = 'g'")
		query, err := buildQuery(query, filter.G)
		if err != nil {
			return err
		}
		err = query.Select()
		if err != nil {
			return err
		}

		for _, line := range lines {
			handler(line.String(), model)
		}
	}
	return nil
}

func (a *Adapter) IsFiltered() bool {
	return a.filtered
}