diff --git a/pkg/webhook/statefulset/statefulset.go b/pkg/webhook/statefulset/statefulset.go
index 26f8b6ed185..e975e3432f2 100644
--- a/pkg/webhook/statefulset/statefulset.go
+++ b/pkg/webhook/statefulset/statefulset.go
@@ -14,12 +14,10 @@
 package statefulset
 
 import (
-	"errors"
 	"fmt"
-	"k8s.io/klog"
 	"strconv"
 
-	asappsv1 "github.com/pingcap/advanced-statefulset/client/apis/apps/v1"
+	asapps "github.com/pingcap/advanced-statefulset/client/apis/apps/v1"
 	"github.com/pingcap/tidb-operator/pkg/client/clientset/versioned"
 	"github.com/pingcap/tidb-operator/pkg/controller"
 	"github.com/pingcap/tidb-operator/pkg/features"
@@ -29,16 +27,13 @@ import (
 	apps "k8s.io/api/apps/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/klog"
 )
 
 var (
-	deserializer runtime.Decoder
+	deserializer runtime.Decoder = util.Codecs.UniversalDeserializer()
 )
 
-func init() {
-	deserializer = util.GetCodec()
-}
-
 type StatefulSetAdmissionControl struct {
 	// operator client interface
 	operatorCli versioned.Interface
@@ -56,14 +51,14 @@ func (sc *StatefulSetAdmissionControl) AdmitStatefulSets(ar *admission.Admission
 	namespace := ar.Namespace
 	expectedGroup := "apps"
 	if features.DefaultFeatureGate.Enabled(features.AdvancedStatefulSet) {
-		expectedGroup = asappsv1.GroupName
+		expectedGroup = asapps.GroupName
 	}
 	apiVersion := ar.Resource.Version
 	setResource := metav1.GroupVersionResource{Group: expectedGroup, Version: apiVersion, Resource: "statefulsets"}
 
 	klog.Infof("admit %s [%s/%s]", setResource, namespace, name)
 
-	stsObjectMeta, stsPartition, err := getStsAttributes(ar.OldObject.Raw)
+	stsObjectMeta, stsPartition, err := getStsAttributes(ar.Object.Raw)
 	if err != nil {
 		err = fmt.Errorf("statefulset %s/%s, decode request failed, err: %v", namespace, name, err)
 		klog.Error(err)
@@ -93,11 +88,11 @@ func (sc *StatefulSetAdmissionControl) AdmitStatefulSets(ar *admission.Admission
 		return util.ARFail(err)
 	}
 
-	var partitionStr string
-	partitionStr = tc.Annotations[label.AnnTiDBPartition]
+	annKey := label.AnnTiDBPartition
 	if l.IsTiKV() {
-		partitionStr = tc.Annotations[label.AnnTiKVPartition]
+		annKey = label.AnnTiKVPartition
 	}
+	partitionStr := tc.Annotations[annKey]
 
 	if len(partitionStr) == 0 {
 		return util.ARSuccess()
@@ -111,9 +106,10 @@ func (sc *StatefulSetAdmissionControl) AdmitStatefulSets(ar *admission.Admission
 	}
 
 	if stsPartition != nil {
-		if *stsPartition > 0 && *stsPartition <= int32(partition) {
-			klog.Infof("statefulset %s/%s has been protect by partition %s annotations", namespace, name, partitionStr)
-			return util.ARFail(errors.New("protect by partition annotation"))
+		// only [partition from tc, INT32_MAX] are allowed
+		if *stsPartition < int32(partition) {
+			klog.Infof("statefulset %s/%s has been protected by partition annotation %q", namespace, name, partitionStr)
+			return util.ARFail(fmt.Errorf("protected by partition annotation (%s = %s) on the tidb cluster %s/%s", annKey, partitionStr, namespace, tcName))
 		}
 		klog.Infof("admit statefulset %s/%s update partition to %d, protect partition is %d", namespace, name, *stsPartition, partition)
 	}
@@ -121,7 +117,17 @@ func (sc *StatefulSetAdmissionControl) AdmitStatefulSets(ar *admission.Admission
 }
 
 func getStsAttributes(data []byte) (*metav1.ObjectMeta, *int32, error) {
-	set := apps.StatefulSet{}
+	if !features.DefaultFeatureGate.Enabled(features.AdvancedStatefulSet) {
+		set := apps.StatefulSet{}
+		if _, _, err := deserializer.Decode(data, nil, &set); err != nil {
+			return nil, nil, err
+		}
+		if set.Spec.UpdateStrategy.RollingUpdate != nil {
+			return &(set.ObjectMeta), set.Spec.UpdateStrategy.RollingUpdate.Partition, nil
+		}
+		return &(set.ObjectMeta), nil, nil
+	}
+	set := asapps.StatefulSet{}
 	if _, _, err := deserializer.Decode(data, nil, &set); err != nil {
 		return nil, nil, err
 	}
diff --git a/pkg/webhook/statefulset/statefulset_test.go b/pkg/webhook/statefulset/statefulset_test.go
new file mode 100644
index 00000000000..62d61831c62
--- /dev/null
+++ b/pkg/webhook/statefulset/statefulset_test.go
@@ -0,0 +1,285 @@
+// Copyright 2020 PingCAP, Inc.
+//
+// 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,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package statefulset
+
+import (
+	"bytes"
+	"testing"
+
+	asapps "github.com/pingcap/advanced-statefulset/client/apis/apps/v1"
+	"github.com/pingcap/advanced-statefulset/client/apis/apps/v1/helper"
+	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
+	"github.com/pingcap/tidb-operator/pkg/client/clientset/versioned/fake"
+	"github.com/pingcap/tidb-operator/pkg/features"
+	"github.com/pingcap/tidb-operator/pkg/webhook/util"
+	admission "k8s.io/api/admission/v1beta1"
+	apps "k8s.io/api/apps/v1"
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/utils/pointer"
+)
+
+type testcase struct {
+	name        string
+	sts         *apps.StatefulSet
+	tc          *v1alpha1.TidbCluster
+	operation   admission.Operation
+	wantAllowed bool
+}
+
+var (
+	ownerTCName    = "foo"
+	validOwnerRefs = []metav1.OwnerReference{
+		{
+			APIVersion: "pingcap.com/v1alpha1",
+			Kind:       "TidbCluster",
+			Name:       ownerTCName,
+			Controller: pointer.BoolPtr(true),
+		},
+	}
+	tests = []testcase{
+		{
+			name:        "non-update operation",
+			operation:   admission.Delete,
+			wantAllowed: true,
+		},
+		{
+			name:      "not tidb or tikv",
+			operation: admission.Update,
+			sts: &apps.StatefulSet{
+				TypeMeta: metav1.TypeMeta{
+					Kind:       "StatefulSet",
+					APIVersion: "apps/v1",
+				},
+				ObjectMeta: metav1.ObjectMeta{},
+			},
+			wantAllowed: true,
+		},
+		{
+			name:      "not owned by tidb cluster",
+			operation: admission.Update,
+			sts: &apps.StatefulSet{
+				TypeMeta: metav1.TypeMeta{
+					Kind:       "StatefulSet",
+					APIVersion: "apps/v1",
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{
+						"app.kubernetes.io/component": "tikv",
+					},
+				},
+			},
+			wantAllowed: true,
+		},
+		{
+			name:      "owned by tidb cluster but the cluster does not exist",
+			operation: admission.Update,
+			sts: &apps.StatefulSet{
+				TypeMeta: metav1.TypeMeta{
+					Kind:       "StatefulSet",
+					APIVersion: "apps/v1",
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{
+						"app.kubernetes.io/component": "tikv",
+					},
+					OwnerReferences: validOwnerRefs,
+				},
+			},
+			wantAllowed: false,
+		},
+		{
+			name:      "partition does not exist",
+			operation: admission.Update,
+			sts: &apps.StatefulSet{
+				TypeMeta: metav1.TypeMeta{
+					Kind:       "StatefulSet",
+					APIVersion: "apps/v1",
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{
+						"app.kubernetes.io/component": "tikv",
+					},
+					OwnerReferences: validOwnerRefs,
+				},
+			},
+			tc: &v1alpha1.TidbCluster{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:        ownerTCName,
+					Namespace:   v1.NamespaceDefault,
+					Annotations: nil,
+				},
+			},
+			wantAllowed: true,
+		},
+		{
+			name:      "partition is invald",
+			operation: admission.Update,
+			sts: &apps.StatefulSet{
+				TypeMeta: metav1.TypeMeta{
+					Kind:       "StatefulSet",
+					APIVersion: "apps/v1",
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{
+						"app.kubernetes.io/component": "tikv",
+					},
+					OwnerReferences: validOwnerRefs,
+				},
+			},
+			tc: &v1alpha1.TidbCluster{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      ownerTCName,
+					Namespace: v1.NamespaceDefault,
+					Annotations: map[string]string{
+						"tidb.pingcap.com/tikv-partition": "invalid",
+					},
+				},
+			},
+			wantAllowed: false,
+		},
+		{
+			name:      "partition allow",
+			operation: admission.Update,
+			sts: &apps.StatefulSet{
+				TypeMeta: metav1.TypeMeta{
+					Kind:       "StatefulSet",
+					APIVersion: "apps/v1",
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{
+						"app.kubernetes.io/component": "tikv",
+					},
+					OwnerReferences: validOwnerRefs,
+				},
+				Spec: apps.StatefulSetSpec{
+					UpdateStrategy: apps.StatefulSetUpdateStrategy{
+						Type: apps.RollingUpdateStatefulSetStrategyType,
+						RollingUpdate: &apps.RollingUpdateStatefulSetStrategy{
+							Partition: pointer.Int32Ptr(1),
+						},
+					},
+				},
+			},
+			tc: &v1alpha1.TidbCluster{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      ownerTCName,
+					Namespace: v1.NamespaceDefault,
+					Annotations: map[string]string{
+						"tidb.pingcap.com/tikv-partition": "1",
+					},
+				},
+			},
+			wantAllowed: true,
+		},
+		{
+			name:      "partition deny",
+			operation: admission.Update,
+			sts: &apps.StatefulSet{
+				TypeMeta: metav1.TypeMeta{
+					Kind:       "StatefulSet",
+					APIVersion: "apps/v1",
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{
+						"app.kubernetes.io/component": "tikv",
+					},
+					OwnerReferences: validOwnerRefs,
+				},
+				Spec: apps.StatefulSetSpec{
+					UpdateStrategy: apps.StatefulSetUpdateStrategy{
+						Type: apps.RollingUpdateStatefulSetStrategyType,
+						RollingUpdate: &apps.RollingUpdateStatefulSetStrategy{
+							Partition: pointer.Int32Ptr(0),
+						},
+					},
+				},
+			},
+			tc: &v1alpha1.TidbCluster{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      ownerTCName,
+					Namespace: v1.NamespaceDefault,
+					Annotations: map[string]string{
+						"tidb.pingcap.com/tikv-partition": "1",
+					},
+				},
+			},
+			wantAllowed: false,
+		},
+	}
+)
+
+func runTest(t *testing.T, tt testcase, asts bool) {
+	jsonInfo, ok := runtime.SerializerInfoForMediaType(util.Codecs.SupportedMediaTypes(), runtime.ContentTypeJSON)
+	if !ok {
+		t.Fatalf("unable to locate encoder -- %q is not a supported media type", runtime.ContentTypeJSON)
+	}
+
+	cli := fake.NewSimpleClientset()
+	ac := NewStatefulSetAdmissionControl(cli)
+	ar := &admission.AdmissionRequest{
+		Name:      "foo",
+		Namespace: v1.NamespaceDefault,
+		Operation: tt.operation,
+	}
+	if tt.sts != nil {
+		buf := bytes.Buffer{}
+		if asts {
+			encoder := util.Codecs.EncoderForVersion(jsonInfo.Serializer, asapps.SchemeGroupVersion)
+			sts, err := helper.FromBuiltinStatefulSet(tt.sts)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := encoder.Encode(sts, &buf); err != nil {
+				t.Fatal(err)
+			}
+		} else {
+			encoder := util.Codecs.EncoderForVersion(jsonInfo.Serializer, apps.SchemeGroupVersion)
+			if err := encoder.Encode(tt.sts, &buf); err != nil {
+				t.Fatal(err)
+			}
+		}
+		ar.Object = runtime.RawExtension{
+			Raw: buf.Bytes(),
+		}
+	}
+	if tt.tc != nil {
+		cli.PingcapV1alpha1().TidbClusters(tt.tc.Namespace).Create(tt.tc)
+	}
+	resp := ac.AdmitStatefulSets(ar)
+	if resp.Allowed != tt.wantAllowed {
+		t.Errorf("want allowed %v, got %v", tt.wantAllowed, resp.Allowed)
+	}
+}
+
+func TestStatefulSetAdmissionControl(t *testing.T) {
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			runTest(t, tt, false)
+		})
+	}
+}
+
+func TestStatefulSetAdmissionControl_ASTS(t *testing.T) {
+	saved := features.DefaultFeatureGate.String()
+	features.DefaultFeatureGate.Set("AdvancedStatefulSet=true")
+	defer features.DefaultFeatureGate.Set(saved) // reset features on exit
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			runTest(t, tt, true)
+		})
+	}
+}
diff --git a/pkg/webhook/util/scheme.go b/pkg/webhook/util/scheme.go
index d655a35940b..aa0e9e6dc7e 100644
--- a/pkg/webhook/util/scheme.go
+++ b/pkg/webhook/util/scheme.go
@@ -14,15 +14,23 @@
 package util
 
 import (
+	asappsv1 "github.com/pingcap/advanced-statefulset/client/apis/apps/v1"
 	admissionv1beta1 "k8s.io/api/admission/v1beta1"
 	admissionregistrationv1beta1 "k8s.io/api/admissionregistration/v1beta1"
+	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/serializer"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 )
 
-var scheme = runtime.NewScheme()
+var (
+	// scheme is the runtime.Scheme to which all api types used by webhook are registered.
+	scheme = runtime.NewScheme()
+
+	// Codecs provides access to encoding and decoding for the scheme.
+	Codecs = serializer.NewCodecFactory(scheme)
+)
 
 func init() {
 	addToScheme(scheme)
@@ -30,10 +38,8 @@ func init() {
 
 func addToScheme(scheme *runtime.Scheme) {
 	utilruntime.Must(corev1.AddToScheme(scheme))
+	utilruntime.Must(appsv1.AddToScheme(scheme))
+	utilruntime.Must(asappsv1.AddToScheme(scheme))
 	utilruntime.Must(admissionv1beta1.AddToScheme(scheme))
 	utilruntime.Must(admissionregistrationv1beta1.AddToScheme(scheme))
 }
-
-func GetCodec() runtime.Decoder {
-	return serializer.NewCodecFactory(scheme).UniversalDeserializer()
-}